diff --git a/.env.template b/.env.template
index 3b7c8e887..ba7f752db 100644
--- a/.env.template
+++ b/.env.template
@@ -23,6 +23,15 @@ WEB_SERVER_PORT=7860
#*******************************************************************#
# LLM_MODEL, see /pilot/configs/model_config.LLM_MODEL_CONFIG
LLM_MODEL=vicuna-13b-v1.5
+## LLM model path, by default, DB-GPT will read the model path from LLM_MODEL_CONFIG based on the LLM_MODEL.
+## Of course you can specify your model path according to LLM_MODEL_PATH
+## In DB-GPT, the priority from high to low to read model path:
+## 1. environment variable with key: {LLM_MODEL}_MODEL_PATH (Avoid multi-model conflicts)
+## 2. environment variable with key: MODEL_PATH
+## 3. environment variable with key: LLM_MODEL_PATH
+## 4. the config in /pilot/configs/model_config.LLM_MODEL_CONFIG
+# LLM_MODEL_PATH=/app/models/vicuna-13b-v1.5
+# LLM_PROMPT_TEMPLATE=vicuna_v1.1
MODEL_SERVER=http://127.0.0.1:8000
LIMIT_MODEL_CONCURRENCY=5
MAX_POSITION_EMBEDDINGS=4096
@@ -46,6 +55,17 @@ QUANTIZE_8bit=True
## Model path
# llama_cpp_model_path=/data/models/TheBloke/vicuna-13B-v1.5-GGUF/vicuna-13b-v1.5.Q4_K_M.gguf
+### LLM cache
+## Enable Model cache
+# MODEL_CACHE_ENABLE=True
+## The storage type of model cache, now supports: memory, disk
+# MODEL_CACHE_STORAGE_TYPE=disk
+## The max cache data in memory, we always store cache data in memory fist for high speed.
+# MODEL_CACHE_MAX_MEMORY_MB=256
+## The dir to save cache data, this configuration is only valid when MODEL_CACHE_STORAGE_TYPE=disk
+## The default dir is pilot/data/model_cache
+# MODEL_CACHE_STORAGE_DISK_DIR=
+
#*******************************************************************#
#** EMBEDDING SETTINGS **#
#*******************************************************************#
@@ -84,6 +104,7 @@ LOCAL_DB_TYPE=sqlite
# LOCAL_DB_PASSWORD=aa12345678
# LOCAL_DB_HOST=127.0.0.1
# LOCAL_DB_PORT=3306
+# LOCAL_DB_NAME=dbgpt
### This option determines the storage location of conversation records. The default is not configured to the old version of duckdb. It can be optionally db or file (if the value is db, the database configured by LOCAL_DB will be used)
#CHAT_HISTORY_STORE_TYPE=db
diff --git a/CODE_OF_CONDUCT b/CODE_OF_CONDUCT
new file mode 100644
index 000000000..b7efcc0b3
--- /dev/null
+++ b/CODE_OF_CONDUCT
@@ -0,0 +1,126 @@
+# Contributor Covenant 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
+
+Community leaders 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.
+
+Community leaders 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 community leaders responsible for enforcement at
+[INSERT CONTACT METHOD].
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+*Community Impact*: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+*Consequence*: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+*Community Impact*: A violation through a single incident or series of
+actions.
+
+*Consequence*: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or permanent
+ban.
+
+### 3. Temporary Ban
+
+*Community Impact*: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+*Consequence*: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+*Community Impact*: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior, harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+*Consequence*: A permanent ban from any sort of public interaction within the
+community.
+
+## 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.html][v2.1].
+
+Community Impact Guidelines were inspired by
+[Mozilla's code of conduct enforcement ladder][Mozilla CoC].
+
+For answers to common questions about this code of conduct, see the FAQ at
+[https://www.contributor-covenant.org/faq][FAQ]. Translations are available at
+[https://www.contributor-covenant.org/translations][translations].
diff --git a/README.md b/README.md
index 49b006a7b..524541209 100644
--- a/README.md
+++ b/README.md
@@ -13,9 +13,6 @@
-
-
-
@@ -25,8 +22,8 @@
-
-
+
+
@@ -34,17 +31,27 @@
-
-
diff --git a/README.zh.md b/README.zh.md
index a6119663d..00ead0cc8 100644
--- a/README.zh.md
+++ b/README.zh.md
@@ -22,24 +22,21 @@
-
-
+
+
-
-
-
@@ -345,16 +323,6 @@ The MIT License (MIT)
- SFT模型准确率
截止20231010,我们利用本项目基于开源的13B大小的模型微调后,在Spider的评估集上的执行准确率,已经超越GPT-4!
-| 模型名称 | 执行准确率 | 说明 |
-| ----------------------------------| ------------------ | ------------------------------------------------------------------------------------------------------------------------------ |
-| **GPT-4** | **0.762** | [numbersstation-eval-res](https://www.numbersstation.ai/post/nsql-llama-2-7b) |
-| ChatGPT | 0.728 | [numbersstation-eval-res](https://www.numbersstation.ai/post/nsql-llama-2-7b) |
-| **CodeLlama-13b-Instruct-hf_lora**| **0.789** | sft train by our this project,only used spider train dataset ,the same eval way in this project with lora SFT |
-| CodeLlama-13b-Instruct-hf_qlora | 0.774 | sft train by our this project,only used spider train dataset ,the same eval way in this project with qlora and nf4,bit4 SFT |
-| wizardcoder | 0.610 | [text-to-sql-wizardcoder](https://github.com/cuplv/text-to-sql-wizardcoder/tree/main) |
-| CodeLlama-13b-Instruct-hf | 0.556 | eval in this project default param |
-| llama2_13b_hf_lora_best | 0.744 | sft train by our this project,only used spider train dataset ,the same eval way in this project |
-
[More Information about Text2SQL finetune](https://github.com/eosphoros-ai/DB-GPT-Hub)
## 联系我们
diff --git a/assets/schema/history.sql b/assets/schema/history.sql
deleted file mode 100644
index 3323b73a0..000000000
--- a/assets/schema/history.sql
+++ /dev/null
@@ -1,18 +0,0 @@
-CREATE DATABASE history;
-use history;
-CREATE TABLE `chat_feed_back` (
- `id` bigint(20) NOT NULL AUTO_INCREMENT,
- `conv_uid` varchar(128) DEFAULT NULL COMMENT '会话id',
- `conv_index` int(4) DEFAULT NULL COMMENT '第几轮会话',
- `score` int(1) DEFAULT NULL COMMENT '评分',
- `ques_type` varchar(32) DEFAULT NULL COMMENT '用户问题类别',
- `question` longtext DEFAULT NULL COMMENT '用户问题',
- `knowledge_space` varchar(128) DEFAULT NULL COMMENT '知识库',
- `messages` longtext DEFAULT NULL COMMENT '评价详情',
- `user_name` varchar(128) DEFAULT NULL COMMENT '评价人',
- `gmt_created` datetime DEFAULT NULL,
- `gmt_modified` datetime DEFAULT NULL,
- PRIMARY KEY (`id`),
- UNIQUE KEY `uk_conv` (`conv_uid`,`conv_index`),
- KEY `idx_conv` (`conv_uid`,`conv_index`)
-) ENGINE=InnoDB AUTO_INCREMENT=0 DEFAULT CHARSET=utf8mb4 COMMENT='用户评分反馈表';
\ No newline at end of file
diff --git a/assets/schema/knowledge_management.sql b/assets/schema/knowledge_management.sql
index 95ed760e6..18f632b14 100644
--- a/assets/schema/knowledge_management.sql
+++ b/assets/schema/knowledge_management.sql
@@ -1,5 +1,15 @@
-CREATE DATABASE knowledge_management;
-use knowledge_management;
+-- You can change `dbgpt` to your actual metadata database name in your `.env` file
+-- eg. `LOCAL_DB_NAME=dbgpt`
+
+CREATE DATABASE IF NOT EXISTS dbgpt;
+use dbgpt;
+
+-- For alembic migration tool
+CREATE TABLE `alembic_version` (
+ version_num VARCHAR(32) NOT NULL,
+ CONSTRAINT alembic_version_pkc PRIMARY KEY (version_num)
+);
+
CREATE TABLE `knowledge_space` (
`id` int NOT NULL AUTO_INCREMENT COMMENT 'auto increment id',
`name` varchar(100) NOT NULL COMMENT 'knowledge space name',
@@ -26,6 +36,7 @@ CREATE TABLE `knowledge_document` (
`content` LONGTEXT NOT NULL COMMENT 'knowledge embedding sync result',
`result` TEXT NULL COMMENT 'knowledge content',
`vector_ids` LONGTEXT NULL COMMENT 'vector_ids',
+ `summary` LONGTEXT NULL COMMENT 'knowledge summary',
`gmt_created` TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT 'created time',
`gmt_modified` TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT 'update time',
PRIMARY KEY (`id`),
@@ -45,6 +56,102 @@ CREATE TABLE `document_chunk` (
KEY `idx_document_id` (`document_id`) COMMENT 'index:document_id'
) ENGINE=InnoDB AUTO_INCREMENT=100001 DEFAULT CHARSET=utf8mb4 COMMENT='knowledge document chunk detail';
+
+
+CREATE TABLE `connect_config` (
+ `id` int NOT NULL AUTO_INCREMENT COMMENT 'autoincrement id',
+ `db_type` varchar(255) NOT NULL COMMENT 'db type',
+ `db_name` varchar(255) NOT NULL COMMENT 'db name',
+ `db_path` varchar(255) DEFAULT NULL COMMENT 'file db path',
+ `db_host` varchar(255) DEFAULT NULL COMMENT 'db connect host(not file db)',
+ `db_port` varchar(255) DEFAULT NULL COMMENT 'db cnnect port(not file db)',
+ `db_user` varchar(255) DEFAULT NULL COMMENT 'db user',
+ `db_pwd` varchar(255) DEFAULT NULL COMMENT 'db password',
+ `comment` text COMMENT 'db comment',
+ PRIMARY KEY (`id`),
+ UNIQUE KEY `uk_db` (`db_name`),
+ KEY `idx_q_db_type` (`db_type`)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT 'Connection confi';
+
+CREATE TABLE `chat_history` (
+ `id` int NOT NULL AUTO_INCREMENT COMMENT 'autoincrement id',
+ `conv_uid` varchar(255) COLLATE utf8mb4_unicode_ci NOT NULL COMMENT 'Conversation record unique id',
+ `chat_mode` varchar(255) COLLATE utf8mb4_unicode_ci NOT NULL COMMENT 'Conversation scene mode',
+ `summary` varchar(255) COLLATE utf8mb4_unicode_ci NOT NULL COMMENT 'Conversation record summary',
+ `user_name` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'interlocutor',
+ `messages` text COLLATE utf8mb4_unicode_ci COMMENT 'Conversation details',
+ PRIMARY KEY (`id`)
+) ENGINE=InnoDB AUTO_INCREMENT=2 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT 'Chat history';
+
+CREATE TABLE `chat_feed_back` (
+ `id` bigint(20) NOT NULL AUTO_INCREMENT,
+ `conv_uid` varchar(128) DEFAULT NULL COMMENT 'Conversation ID',
+ `conv_index` int(4) DEFAULT NULL COMMENT 'Round of conversation',
+ `score` int(1) DEFAULT NULL COMMENT 'Score of user',
+ `ques_type` varchar(32) DEFAULT NULL COMMENT 'User question category',
+ `question` longtext DEFAULT NULL COMMENT 'User question',
+ `knowledge_space` varchar(128) DEFAULT NULL COMMENT 'Knowledge space name',
+ `messages` longtext DEFAULT NULL COMMENT 'The details of user feedback',
+ `user_name` varchar(128) DEFAULT NULL COMMENT 'User name',
+ `gmt_created` timestamp NULL DEFAULT CURRENT_TIMESTAMP COMMENT 'created time',
+ `gmt_modified` timestamp NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT 'update time',
+ PRIMARY KEY (`id`),
+ UNIQUE KEY `uk_conv` (`conv_uid`,`conv_index`),
+ KEY `idx_conv` (`conv_uid`,`conv_index`)
+) ENGINE=InnoDB AUTO_INCREMENT=0 DEFAULT CHARSET=utf8mb4 COMMENT='User feedback table';
+
+
+CREATE TABLE `my_plugin` (
+ `id` int NOT NULL AUTO_INCREMENT COMMENT 'autoincrement id',
+ `tenant` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'user tenant',
+ `user_code` varchar(255) COLLATE utf8mb4_unicode_ci NOT NULL COMMENT 'user code',
+ `user_name` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'user name',
+ `name` varchar(255) COLLATE utf8mb4_unicode_ci NOT NULL COMMENT 'plugin name',
+ `file_name` varchar(255) COLLATE utf8mb4_unicode_ci NOT NULL COMMENT 'plugin package file name',
+ `type` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'plugin type',
+ `version` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'plugin version',
+ `use_count` int DEFAULT NULL COMMENT 'plugin total use count',
+ `succ_count` int DEFAULT NULL COMMENT 'plugin total success count',
+ `gmt_created` TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT 'plugin install time',
+ PRIMARY KEY (`id`),
+ UNIQUE KEY `name` (`name`)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='User plugin table';
+
+CREATE TABLE `plugin_hub` (
+ `id` int NOT NULL AUTO_INCREMENT COMMENT 'autoincrement id',
+ `name` varchar(255) COLLATE utf8mb4_unicode_ci NOT NULL COMMENT 'plugin name',
+ `description` varchar(255) COLLATE utf8mb4_unicode_ci NOT NULL COMMENT 'plugin description',
+ `author` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'plugin author',
+ `email` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'plugin author email',
+ `type` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'plugin type',
+ `version` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'plugin version',
+ `storage_channel` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'plugin storage channel',
+ `storage_url` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'plugin download url',
+ `download_param` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'plugin download param',
+ `gmt_created` TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT 'plugin upload time',
+ `installed` int DEFAULT NULL COMMENT 'plugin already installed count',
+ PRIMARY KEY (`id`),
+ UNIQUE KEY `name` (`name`)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='Plugin Hub table';
+
+
+CREATE TABLE `prompt_manage` (
+ `id` int(11) NOT NULL AUTO_INCREMENT,
+ `chat_scene` varchar(100) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'Chat scene',
+ `sub_chat_scene` varchar(100) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'Sub chat scene',
+ `prompt_type` varchar(100) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'Prompt type: common or private',
+ `prompt_name` varchar(512) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'prompt name',
+ `content` longtext COLLATE utf8mb4_unicode_ci COMMENT 'Prompt content',
+ `user_name` varchar(128) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'User name',
+ `gmt_created` timestamp NULL DEFAULT CURRENT_TIMESTAMP COMMENT 'created time',
+ `gmt_modified` timestamp NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT 'update time',
+ PRIMARY KEY (`id`),
+ UNIQUE KEY `prompt_name_uiq` (`prompt_name`),
+ KEY `gmt_created_idx` (`gmt_created`)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='Prompt management table';
+
+
+
CREATE DATABASE EXAMPLE_1;
use EXAMPLE_1;
CREATE TABLE `users` (
diff --git a/assets/schema/prompt_management.sql b/assets/schema/prompt_management.sql
deleted file mode 100644
index b2ed6de23..000000000
--- a/assets/schema/prompt_management.sql
+++ /dev/null
@@ -1,16 +0,0 @@
-CREATE DATABASE prompt_management;
-use prompt_management;
-CREATE TABLE `prompt_manage` (
- `id` int(11) NOT NULL AUTO_INCREMENT,
- `chat_scene` varchar(100) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT '场景',
- `sub_chat_scene` varchar(100) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT '子场景',
- `prompt_type` varchar(100) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT '类型: common or private',
- `prompt_name` varchar(512) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'prompt的名字',
- `content` longtext COLLATE utf8mb4_unicode_ci COMMENT 'prompt的内容',
- `user_name` varchar(128) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT '用户名',
- `gmt_created` datetime DEFAULT NULL,
- `gmt_modified` datetime DEFAULT NULL,
- PRIMARY KEY (`id`),
- UNIQUE KEY `prompt_name_uiq` (`prompt_name`),
- KEY `gmt_created_idx` (`gmt_created`)
-) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='prompt管理表';
\ No newline at end of file
diff --git a/assets/wechat.jpg b/assets/wechat.jpg
index 641de1489..7a2ce6dd5 100644
Binary files a/assets/wechat.jpg and b/assets/wechat.jpg differ
diff --git a/docker/base/Dockerfile b/docker/base/Dockerfile
index 63486e260..e4fad9697 100644
--- a/docker/base/Dockerfile
+++ b/docker/base/Dockerfile
@@ -28,7 +28,9 @@ WORKDIR /app
# RUN pip3 install -i $PIP_INDEX_URL ".[all]"
RUN pip3 install --upgrade pip -i $PIP_INDEX_URL \
- && pip3 install -i $PIP_INDEX_URL ".[$DB_GPT_INSTALL_MODEL]"
+ && pip3 install -i $PIP_INDEX_URL ".[$DB_GPT_INSTALL_MODEL]" \
+ # install openai for proxyllm
+ && pip3 install -i $PIP_INDEX_URL ".[openai]"
RUN (if [ "${LANGUAGE}" = "zh" ]; \
# language is zh, download zh_core_web_sm from github
diff --git a/docker/compose_examples/cluster-docker-compose.yml b/docker/compose_examples/cluster-docker-compose.yml
index b41033458..0ad6be9ae 100644
--- a/docker/compose_examples/cluster-docker-compose.yml
+++ b/docker/compose_examples/cluster-docker-compose.yml
@@ -7,6 +7,16 @@ services:
restart: unless-stopped
networks:
- dbgptnet
+ api-server:
+ image: eosphorosai/dbgpt:latest
+ command: dbgpt start apiserver --controller_addr http://controller:8000
+ restart: unless-stopped
+ depends_on:
+ - controller
+ networks:
+ - dbgptnet
+ ports:
+ - 8100:8100/tcp
llm-worker:
image: eosphorosai/dbgpt:latest
command: dbgpt start worker --model_name vicuna-13b-v1.5 --model_path /app/models/vicuna-13b-v1.5 --port 8001 --controller_addr http://controller:8000
diff --git a/pilot/datasets/oceanbase/OceanBase_Introduction.md b/docker/examples/benchmarks/benchmarks_llm_11k_prompt.txt
similarity index 97%
rename from pilot/datasets/oceanbase/OceanBase_Introduction.md
rename to docker/examples/benchmarks/benchmarks_llm_11k_prompt.txt
index 2cb6e9538..25072fa18 100644
--- a/pilot/datasets/oceanbase/OceanBase_Introduction.md
+++ b/docker/examples/benchmarks/benchmarks_llm_11k_prompt.txt
@@ -1,3 +1,8 @@
+A chat between a curious user and an artificial intelligence assistant, who very familiar with database related knowledge.
+The assistant gives helpful, detailed, professional and polite answers to the user's questions. 基于以下已知的信息, 专业、简要的回答用户的问题,
+ 如果无法从提供的内容中获取答案, 请说: "知识库中提供的内容不足以回答此问题" 禁止胡乱编造, 回答的时候最好按照1.2.3.点进行总结。
+ 已知内容:
+
OceanBase 数据库(OceanBase Database)是一款完全自研的企业级原生分布式数据库,在普通硬件上实现金融级高可用,首创“三地五中心”城市级故障自动无损容灾新标准,刷新 TPC-C 标准测试,单集群规模超过 1500 节点,具有云原生、强一致性、高度兼容 Oracle/MySQL 等特性。
核心特性
@@ -203,4 +208,7 @@ OceanBase 数据库是多租户的数据库系统,一个集群内可包含多
创建租户前,需首先确定租户的资源配置、使用资源范围等。租户创建的通用流程如下:
-资源配置是描述资源池的配置信息,用来描述资源池中每个资源单元可用的 CPU、内存、存储空间和 IOPS 等的规格。修改资源配置可动态调整资源单元的规格。这里需要注意,资源配置指定的是对应资源单元能够提供的服务能力,而不是资源单元的实时负载。 创建资源配置的示例语句如下:
\ No newline at end of file
+资源配置是描述资源池的配置信息,用来描述资源池中每个资源单元可用的 CPU、内存、存储空间和 IOPS 等的规格。修改资源配置可动态调整资源单元的规格。这里需要注意,资源配置指定的是对应资源单元能够提供的服务能力,而不是资源单元的实时负载。 创建资源配置的示例语句如下:
+
+问题:
+请你基于上述内容对 OceanBase 的介绍进行总结,不少于2000字。
\ No newline at end of file
diff --git a/docker/examples/dashboard/test_case_mysql_data.py b/docker/examples/dashboard/test_case_mysql_data.py
new file mode 100644
index 000000000..e2ef0badd
--- /dev/null
+++ b/docker/examples/dashboard/test_case_mysql_data.py
@@ -0,0 +1,219 @@
+import random
+import string
+import os
+import pymysql
+from typing import List
+
+import pymysql.cursors
+from datetime import datetime, timedelta
+
+# At first you need to create an test database which called dbgpt_test;
+# you can use next command to create.
+# CREATE DATABASE IF NOT EXISTS dbgpt_test CHARACTER SET utf8;
+
+
+def build_table(connection):
+ connection.cursor().execute(
+ """CREATE TABLE user (
+ id INT(11) NOT NULL AUTO_INCREMENT COMMENT '用户ID',
+ name VARCHAR(50) NOT NULL COMMENT '用户名',
+ email VARCHAR(50) NOT NULL COMMENT '电子邮件',
+ mobile CHAR(11) NOT NULL COMMENT '手机号码',
+ gender VARCHAR(20) COMMENT '性别',
+ birth DATE COMMENT '出生日期',
+ country VARCHAR(20) COMMENT '国家',
+ city VARCHAR(20) COMMENT '城市',
+ create_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
+ update_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
+ PRIMARY KEY (id),
+ UNIQUE KEY uk_email (email)
+ ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='用户信息表';"""
+ )
+ connection.cursor().execute(
+ """CREATE TABLE transaction_order (
+ id INT(11) NOT NULL AUTO_INCREMENT COMMENT '订单ID',
+ order_no CHAR(20) NOT NULL COMMENT '订单编号',
+ product_name VARCHAR(50) NOT NULL COMMENT '产品名称',
+ product_category VARCHAR(20) COMMENT '产品分类',
+ amount DECIMAL(10, 2) NOT NULL COMMENT '订单金额',
+ pay_status VARCHAR(20) COMMENT '付款状态',
+ user_id INT(11) NOT NULL COMMENT '用户ID',
+ user_name VARCHAR(50) COMMENT '用户名',
+ create_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
+ update_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
+ PRIMARY KEY (id),
+ UNIQUE KEY uk_order_no (order_no),
+ KEY idx_user_id (user_id)
+ ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='交易订单表';"""
+ )
+
+
+def user_build(names: List, country: str, grander: str = "Male") -> List:
+ countries = ["China", "US", "India", "Indonesia", "Pakistan"] # 国家列表
+ cities = {
+ "China": ["Beijing", "Shanghai", "Guangzhou", "Shenzhen", "Hangzhou"],
+ "US": ["New York", "Los Angeles", "Chicago", "Houston", "Phoenix"],
+ "India": ["Mumbai", "Delhi", "Bangalore", "Hyderabad", "Chennai"],
+ "Indonesia": ["Jakarta", "Surabaya", "Medan", "Bandung", "Makassar"],
+ "Pakistan": ["Karachi", "Lahore", "Faisalabad", "Rawalpindi", "Multan"],
+ }
+
+ users = []
+ for i in range(1, len(names) + 1):
+ if grander == "Male":
+ id = int(str(countries.index(country) + 1) + "10") + i
+ else:
+ id = int(str(countries.index(country) + 1) + "20") + i
+
+ name = names[i - 1]
+ email = f"{name}@example.com"
+ mobile = "".join(random.choices(string.digits, k=10))
+ gender = grander
+ birth = f"19{random.randint(60, 99)}-{random.randint(1, 12):02d}-{random.randint(1, 28):02d}"
+ country = country
+ city = random.choice(cities[country])
+
+ now = datetime.now()
+ year = now.year
+
+ start = datetime(year, 1, 1)
+ end = datetime(year, 12, 31)
+ random_date = start + timedelta(days=random.randint(0, (end - start).days))
+ random_time = datetime.combine(random_date, datetime.min.time()) + timedelta(
+ seconds=random.randint(0, 24 * 60 * 60 - 1)
+ )
+
+ random_datetime_str = random_time.strftime("%Y-%m-%d %H:%M:%S")
+ create_time = random_datetime_str
+ users.append(
+ (
+ id,
+ name,
+ email,
+ mobile,
+ gender,
+ birth,
+ country,
+ city,
+ create_time,
+ create_time,
+ )
+ )
+ return users
+
+
+def gnerate_all_users(cursor):
+ users = []
+ users_f = ["ZhangWei", "LiQiang", "ZhangSan", "LiSi"]
+ users.extend(user_build(users_f, "China", "Male"))
+ users_m = ["Hanmeimei", "LiMeiMei", "LiNa", "ZhangLi", "ZhangMing"]
+ users.extend(user_build(users_m, "China", "Female"))
+
+ users1_f = ["James", "John", "David", "Richard"]
+ users.extend(user_build(users1_f, "US", "Male"))
+ users1_m = ["Mary", "Patricia", "Sarah"]
+ users.extend(user_build(users1_m, "US", "Female"))
+
+ users2_f = ["Ravi", "Rajesh", "Ajay", "Arjun", "Sanjay"]
+ users.extend(user_build(users2_f, "India", "Male"))
+ users2_m = ["Priya", "Sushma", "Pooja", "Swati"]
+ users.extend(user_build(users2_m, "India", "Female"))
+ for user in users:
+ cursor.execute(
+ "INSERT INTO user (id, name, email, mobile, gender, birth, country, city, create_time, update_time) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)",
+ user,
+ )
+
+ return users
+
+
+def gnerate_all_orders(users, cursor):
+ orders = []
+ orders_num = 200
+ categories = ["Clothing", "Food", "Home Appliance", "Mother and Baby", "Travel"]
+
+ categories_product = {
+ "Clothing": ["T-shirt", "Jeans", "Skirt", "Other"],
+ "Food": ["Snack", "Fruit"],
+ "Home Appliance": ["Refrigerator", "Television", "Air conditioner"],
+ "Mother and Baby": ["Diapers", "Milk Powder", "Stroller", "Toy"],
+ "Travel": ["Tent", "Fishing Rod", "Bike", "Rawalpindi", "Multan"],
+ }
+
+ for i in range(1, orders_num + 1):
+ id = i
+ order_no = "".join(random.choices(string.ascii_uppercase, k=3)) + "".join(
+ random.choices(string.digits, k=10)
+ )
+ product_category = random.choice(categories)
+ product_name = random.choice(categories_product[product_category])
+ amount = round(random.uniform(0, 10000), 2)
+ pay_status = random.choice(["SUCCESS", "FAILD", "CANCEL", "REFUND"])
+ user_id = random.choice(users)[0]
+ user_name = random.choice(users)[1]
+
+ now = datetime.now()
+ year = now.year
+
+ start = datetime(year, 1, 1)
+ end = datetime(year, 12, 31)
+ random_date = start + timedelta(days=random.randint(0, (end - start).days))
+ random_time = datetime.combine(random_date, datetime.min.time()) + timedelta(
+ seconds=random.randint(0, 24 * 60 * 60 - 1)
+ )
+
+ random_datetime_str = random_time.strftime("%Y-%m-%d %H:%M:%S")
+ create_time = random_datetime_str
+
+ order = (
+ id,
+ order_no,
+ product_category,
+ product_name,
+ amount,
+ pay_status,
+ user_id,
+ user_name,
+ create_time,
+ )
+ cursor.execute(
+ "INSERT INTO transaction_order (id, order_no, product_name, product_category, amount, pay_status, user_id, user_name, create_time) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)",
+ order,
+ )
+
+
+if __name__ == "__main__":
+ connection = pymysql.connect(
+ host=os.getenv("DB_HOST", "127.0.0.1"),
+ port=int(
+ os.getenv("DB_PORT", 3306),
+ ),
+ user=os.getenv("DB_USER", "root"),
+ password=os.getenv("DB_PASSWORD", "aa12345678"),
+ database=os.getenv("DB_DATABASE", "dbgpt_test"),
+ charset="utf8mb4",
+ ssl_ca=None,
+ )
+
+ build_table(connection)
+
+ connection.commit()
+
+ cursor = connection.cursor()
+
+ users = gnerate_all_users(cursor)
+ connection.commit()
+
+ gnerate_all_orders(users, cursor)
+ connection.commit()
+
+ cursor.execute("SELECT * FROM user")
+ data = cursor.fetchall()
+ print(data)
+
+ cursor.execute("SELECT count(*) FROM transaction_order")
+ data = cursor.fetchall()
+ print("orders:" + str(data))
+
+ cursor.close()
+ connection.close()
diff --git a/docs/_static/css/custom.css b/docs/_static/css/custom.css
new file mode 100644
index 000000000..f3bcfe7b5
--- /dev/null
+++ b/docs/_static/css/custom.css
@@ -0,0 +1,928 @@
+/* override default colors used in the Sphinx theme */
+:root {
+ --tabs-color-label-active: #0475DE;
+ --tabs-color-label-hover: #0475DE;
+ --buttons-color-blue: #0475DE;
+ --tabs-color-label-inactive: #9E9E9E;
+ --tabs-color-overline: #e0e0e0;
+ --tabs-color-underline: #e0e0e0;
+ --border-color-gray: #e0e0e0;
+ --background-color-light-gray:#fafafa;
+ --background-color-disabled: #9E9E9E;
+ --pst-color-link: 4, 117, 222;
+ --pst-color-primary: 4, 117, 222;
+ --pst-color-text-secondary: #616161;
+ --blue: #0475DE;
+ --sidebar-top: 5em;
+}
+
+/* Remove flicker for announcement top bar replacement */
+.header-item.announcement {
+ background-color: white;
+ color: white;
+ padding: 0;
+}
+
+/* Make the book theme secondary nav stick below the new main top nav */
+.header-article {
+ top: 58px;
+ z-index: 900 !important;
+}
+
+.toctree-l1.has-children {
+ font-weight: bold;
+}
+
+.toctree-l2 {
+ font-weight: normal;
+}
+
+div.navbar-brand-box {
+ padding-top: 4em;
+}
+
+td p {
+ margin-left: 0.75rem;
+}
+
+table.longtable.table.autosummary {
+ table-layout: fixed;
+}
+
+.table.autosummary td {
+ width: 100%;
+}
+
+tr.row-odd {
+ background-color: #f9fafb;
+}
+
+/* For Algolia search box
+ * overflow-y: to flow-over horizontally into main content
+ * height: to prevent topbar overlap
+*/
+#site-navigation {
+ overflow-y: auto;
+ height: calc(100vh - var(--sidebar-top));
+ position: sticky;
+ top: var(--sidebar-top) !important;
+}
+
+/* Center the algolia search bar*/
+#search-input {
+ text-align: center;
+}
+.algolia-autocomplete {
+ width: 100%;
+ margin: auto;
+}
+
+/* Hide confusing "<-" back arrow in navigation for larger displays */
+@media (min-width: 768px) {
+ #navbar-toggler {
+ display: none;
+ }
+}
+
+/* Make navigation scrollable on mobile, by making algolia not overflow */
+@media (max-width: 768px) {
+ #site-navigation {
+ overflow-y: scroll;
+ }
+
+ .algolia-autocomplete .ds-dropdown-menu{
+ min-width: 250px;
+ }
+}
+
+/* sphinx-panels overrides the content width to 1140 for large displays.*/
+@media (min-width: 1200px) {
+ .container, .container-lg, .container-md, .container-sm, .container-xl {
+ max-width: 1400px !important;
+ }
+}
+
+.bottom-right-promo-banner {
+ position: fixed;
+ bottom: 100px;
+ right: 20px;
+ width: 270px;
+}
+
+@media (max-width: 1500px) {
+ .bottom-right-promo-banner {
+ display: none;
+ }
+}
+
+@media screen and (max-width: 767px) {
+ .remove-mobile {
+ display: none;
+ }
+ }
+
+ @media screen and (max-width: 767px) {
+ .row-2-column {
+ flex-direction: column;
+ margin-top: 20px;
+ }
+ }
+
+/* Make Algolia search box scrollable */
+.algolia-autocomplete .ds-dropdown-menu {
+ height: 60vh !important;
+ overflow-y: scroll !important;
+}
+
+.bd-sidebar__content {
+ overflow-y: unset !important;
+}
+
+.bd-sidebar__top {
+ display: flex;
+ flex-direction: column;
+}
+
+.bd-sidebar li {
+ position: relative;
+ word-wrap: break-word;
+}
+
+nav.bd-links {
+ flex: 1;
+}
+
+nav.bd-links::-webkit-scrollbar-thumb {
+ background-color: #ccc;
+}
+
+nav.bd-links::-webkit-scrollbar {
+ width: 5px;
+}
+
+dt:target, span.highlighted {
+ background-color: white;
+}
+
+div.sphx-glr-bigcontainer {
+ display: inline-block;
+ width: 100%;
+}
+
+td.tune-colab,
+th.tune-colab {
+ border: 1px solid #dddddd;
+ text-align: left;
+ padding: 8px;
+}
+
+/* Adjustment to Sphinx Book Theme */
+.table td {
+ /* Remove row spacing on the left */
+ padding-left: 0;
+}
+
+.table thead th {
+ /* Remove row spacing on the left */
+ padding-left: 0;
+}
+
+img.inline-figure {
+ /* Override the display: block for img */
+ display: inherit !important;
+}
+
+#version-warning-banner {
+ /* Make version warning clickable */
+ z-index: 1;
+ margin-left: 0;
+ /* 20% is for ToC rightbar */
+ /* 2 * 1.5625em is for horizontal margins */
+ width: calc(100% - 20% - 2 * 1.5625em);
+}
+
+/* allow scrollable images */
+.figure {
+ max-width: 100%;
+ overflow-x: auto;
+}
+img.horizontal-scroll {
+ max-width: none;
+}
+
+.clear-both {
+ clear: both;
+ min-height: 100px;
+ margin-top: 15px;
+}
+
+.buttons-float-left {
+ width: 150px;
+ float: left;
+}
+
+.buttons-float-right {
+ width: 150px;
+ float: right;
+}
+
+.card-body {
+ padding: 0.5rem !important;
+}
+
+/* custom css for pre elements */
+pre {
+ /* Wrap code blocks instead of horizontal scrolling. */
+ white-space: pre-wrap;
+ box-shadow: none;
+ border-color: var(--border-color-gray);
+ background-color: var(--background-color-light-gray);
+ border-radius:0.25em;
+}
+
+/* notebook formatting */
+.cell .cell_output {
+ max-height: 250px;
+ overflow-y: auto;
+ font-weight: bold;
+}
+
+/* Yellow doesn't render well on light background */
+.cell .cell_output pre .-Color-Yellow {
+ color: #785840;
+}
+
+/* Newlines (\a) and spaces (\20) before each parameter */
+.sig-param::before {
+ content: "\a\20\20\20\20";
+ white-space: pre;
+}
+
+/* custom css for outlined buttons */
+.btn-outline-info:hover span, .btn-outline-primary:hover span {
+ color: #fff;
+}
+
+.btn-outline-info, .btn-outline-primary{
+ border-color: var(--buttons-color-blue);
+}
+
+.btn-outline-info:hover, .btn-outline-primary:hover{
+ border-color: var(--buttons-color-blue);
+ background-color: var(--buttons-color-blue);
+}
+
+.btn-outline-info.active:not(:disabled):not(.disabled), .btn-outline-info:not(:disabled):not(.disabled):active, .show>.btn-outline-info.dropdown-toggle {
+ border-color: var(--buttons-color-blue);
+ background-color: var(--buttons-color-blue);
+ color: #fff;
+}
+
+.btn-info, .btn-info:hover, .btn-info:focus {
+ border-color: var(--buttons-color-blue);
+ background-color: var(--buttons-color-blue);
+}
+
+.btn-info:hover{
+ opacity: 90%;
+}
+
+.btn-info:disabled{
+ border-color: var(--background-color-disabled);
+ background-color: var(--background-color-disabled);
+ opacity: 100%;
+}
+
+.btn-info.active:not(:disabled):not(.disabled), .btn-info:not(:disabled):not(.disabled):active, .show>.btn-info.dropdown-toggle {
+ border-color: var(--buttons-color-blue);
+ background-color: var(--buttons-color-blue);
+}
+
+
+.topnav {
+ background-color: white;
+ border-bottom: 1px solid rgba(0, 0, 0, .1);
+ display: flex;
+ align-items: center;
+}
+
+/* Content wrapper for the unified nav link / menus */
+.top-nav-content {
+ max-width: 1400px;
+ width: 100%;
+ margin-left: auto;
+ margin-right: auto;
+ padding: 0 1.5rem;
+ display: flex;
+ align-items: center;
+ justify-content: space-between;
+}
+
+@media (max-width: 900px) {
+ /* If the window is too small, hide the custom sticky navigation bar at the top of the page.
+ Also make the pydata-sphinx-theme nav bar, which usually sits below the top nav bar, stick
+ to the top of the page.
+ */
+ .top-nav-content {
+ display: none;
+ }
+ div.header-article.row.sticky-top.noprint {
+ position: sticky;
+ top: 0;
+ }
+}
+
+/* Styling the links and menus in the top nav */
+.top-nav-content a {
+ text-decoration: none;
+ color: black;
+ font-size: 17px;
+}
+
+.top-nav-content a:hover {
+ color: #007bff;
+}
+
+/* The left part are the links and menus */
+.top-nav-content > .left {
+ display: flex;
+ white-space: nowrap;
+}
+
+.top-nav-content .left > * {
+ margin-right: 8px;
+}
+
+.top-nav-content .left > a,
+.top-nav-content .left > .menu > a {
+ text-align: center;
+ padding: 14px 16px;
+ border-bottom: 2px solid white;
+}
+
+.top-nav-content .menu:hover > a,
+.top-nav-content .left > a:hover {
+ border-bottom: 2px solid #007bff;
+}
+
+/* Special styling for the Ray logo */
+.top-nav-content .left > a.ray-logo {
+ width: 90px;
+ padding: 10px 0;
+}
+.top-nav-content .left > a.ray-logo:hover {
+ border-bottom: 2px solid white;
+}
+
+/* Styling the dropdown menus */
+.top-nav-content .menu {
+ display: flex;
+}
+.top-nav-content .menu > a > .down-caret {
+ margin-left: 8px;
+}
+.top-nav-content .menu > ul {
+ display: none;
+}
+
+.top-nav-content > button.try-anyscale > span {
+ margin: 0 12px;
+}
+
+.top-nav-content .menu:hover > ul {
+ display: flex;
+ flex-direction: column;
+ align-items: flex-start;
+ box-shadow: 0 5px 15px 0 rgb(0 0 0 / 10%);
+ padding: 15px;
+ width: 330px;
+ position: absolute;
+ z-index: 2000;
+ background-color: white;
+ top: 58px;
+}
+
+.top-nav-content .menu:hover > ul > li {
+ list-style: none;
+ padding: 5px 0;
+}
+
+.top-nav-content .menu:hover > ul > li span {
+ display: block;
+}
+
+.top-nav-content .menu:hover > ul > li span.secondary {
+ color: #787878;
+}
+
+/* Styling the "Try Anyscale" button */
+.top-nav-content > button.try-anyscale {
+ float: right;
+ border-radius: 6px;
+ background-color: #e7f2fa;
+ padding-left: 12px;
+ padding-right: 12px;
+ margin-left: 12px;
+ height: 40px;
+ border: none;
+ white-space: nowrap;
+}
+
+@media (max-width: 1000px) {
+ .top-nav-content > button.try-anyscale {
+ display: none;
+ }
+}
+
+/* custom css for tabs*/
+.tabbed-set>label,.tabbed-set>label:hover {
+ border-bottom: 1px solid var(--border-color-gray);
+ color:var(--tabs-color-label-inactive);
+ font-weight: 500;
+}
+
+.tabbed-set>input:checked+label{
+ border-bottom: 0.125em solid;
+ color:var(--tabs-color-label-active);
+}
+
+
+.tabbed-label{
+ margin-bottom:0;
+}
+
+/* custom css for jupyter cells */
+div.cell div.cell_input{
+ border: 1px var(--border-color-gray) solid;
+ background-color: var(--background-color-light-gray);
+ border-radius:0.25em;
+ border-left-color: var(--green);
+ border-left-width: medium;
+}
+
+/* custom css for table */
+table {
+ border-color: var(--border-color-gray);
+}
+
+/* custom css for topic component */
+div.topic{
+ border: 1px solid var(--border-color-gray);
+ border-radius:0.25em;
+}
+
+.topic {
+ background-color: var(--background-color-light-gray);
+}
+
+/* custom css for card component */
+.card{
+ border-color: var(--border-color-gray);
+}
+
+.card-footer{
+ background-color: var(--background-color-light-gray);
+ border-top-color: var(--border-color-gray);
+}
+
+/* custom css for section navigation component */
+.bd-toc nav>.nav {
+ border-left-color: var(--border-color-gray);
+}
+
+/* custom css for up and down arrows in collapsible cards */
+details.dropdown .summary-up, details.dropdown .summary-down {
+ top: 1em;
+}
+
+/* remove focus border in collapsible admonition buttons */
+.toggle.admonition button.toggle-button:focus {
+ outline: none;
+}
+
+/* custom css for shadow class */
+.shadow {
+ box-shadow: 0 0.2rem 0.5rem rgb(0 0 0 / 5%), 0 0 0.0625rem rgb(0 0 0 / 10%) !important;
+}
+
+/* custom css for text area */
+textarea {
+ border-color: var(--border-color-gray);
+}
+
+/* custom css for footer */
+footer {
+ margin-top: 1rem;
+ padding:1em 0;
+ border-top-color: var(--border-color-gray);
+}
+
+.footer p{
+ color: var(--pst-color-text-secondary);
+}
+
+/* Make the hover color of tag/gallery buttons differ from "active" */
+.tag.btn-outline-primary:hover {
+ background-color: rgba(20, 99, 208, 0.62) !important;
+}
+
+span.rst-current-version > span.fa.fa-book {
+ /* Move the book icon away from the top right
+ * corner of the version flyout menu */
+ margin: 10px 0px 0px 5px;
+}
+
+
+/*Extends the docstring signature box.*/
+.rst-content dl:not(.docutils) dt {
+ display: block;
+ padding: 10px;
+ word-wrap: break-word;
+ padding-right: 100px;
+}
+
+/*Lists in an admonition note do not have awkward whitespace below.*/
+.rst-content .admonition-note .section ul {
+ margin-bottom: 0;
+}
+
+/*Properties become blue (classmethod, staticmethod, property)*/
+.rst-content dl dt em.property {
+ color: #2980b9;
+ text-transform: uppercase;
+}
+
+.rst-content .section ol p,
+.rst-content .section ul p {
+ margin-bottom: 0;
+}
+
+
+/* Adjustment to Version block */
+.rst-versions {
+ z-index: 1200 !important;
+}
+
+.image-header {
+ display: flex;
+ flex-direction: row;
+ align-items: center;
+ padding-left: 16px;
+ padding-right:16px;
+ gap: 16px;
+}
+
+.info-box {
+ box-shadow: 0px 4px 20px rgba(0, 0, 0, 0.05);
+ border-radius: 8px;
+ padding: 20px;
+}
+
+.info-box:hover{
+ box-shadow: 0px 4px 20px rgba(0, 0, 0, 0.1);
+}
+
+.no-underline{
+ text-decoration: none;
+}
+.no-underline:hover{
+ text-decoration: none;
+}
+
+.icon-hover:hover{
+ height: 30px ;
+ width: 30px;
+}
+
+.info-box-2 {
+ background-color: #F9FAFB;
+ border-radius: 8px;
+ padding-right: 16px;
+ padding-left: 16px;
+ padding-bottom: 24px;
+ padding-top: 4px;
+}
+
+
+.bold-link {
+ color: #000000 !important;
+ font-weight: 600;
+}
+
+.community-box {
+ border: 1px solid #D2DCE6;
+ border-radius: 8px;
+ display: flex;
+ margin-bottom: 16px;
+}
+
+.community-box:hover {
+ box-shadow: 0px 4px 20px rgba(0, 0, 0, 0.05);
+ text-decoration: none;
+}
+
+.community-box p {
+ margin-top: 1rem !important;
+}
+
+.tab-pane pre {
+ margin: 0;
+ padding: 0;
+ max-height: 252px;
+ overflow-y: auto;
+}
+
+.grid-container {
+ display: grid;
+ grid-template-columns: repeat(auto-fit, minmax(300px,1fr));
+ grid-gap: 16px;
+}
+
+.grid-item {
+padding: 20px;
+}
+
+
+.nav-pills {
+ background-color: #F9FAFB;
+ color: #000000;
+ padding: 8px;
+ border-bottom:none;
+ border-radius: 8px;
+}
+
+.nav-pills .nav-link.active {
+ background-color: #FFFFFF !important;
+ box-shadow: 0px 3px 14px 2px rgba(3,28,74,0.12);
+ border-radius: 8px;
+ padding: 20px;
+ color: #000000;
+ font-weight: 500;
+}
+
+.searchDiv {
+ width: 100%;
+ position: relative;
+ display: block;
+}
+
+.searchTerm {
+ width: 80%;
+ border: 2px solid var(--blue);
+ padding: 5px;
+ height: 45px;
+ border-radius: 5px;
+ outline: none;
+}
+
+.searchButton {
+ width: 40px;
+ height: 45px;
+ border: 1px solid var(--blue);
+ background: var(--blue);
+ color: #fff;
+ border-radius: 5px;
+ cursor: pointer;
+ font-size: 20px;
+}
+
+/*Resize the wrap to see the search bar change!*/
+.searchWrap {
+ width: 100%;
+ position: relative;
+ margin: 15px;
+ top: 50%;
+ left: 50%;
+ transform: translate(-50%, -10%);
+ text-align: center;
+}
+
+.sd-card {
+ border: none !important;
+}
+
+.tag {
+ margin-bottom: 5px;
+ font-size: small;
+}
+
+/* Override float positioning of next-prev buttons so that
+ they take up space normally, and we can put other stuff at
+ the bottom of the page. */
+.prev-next-area {
+ display: flex;
+ flex-direction: row;
+}
+.prev-next-area a.left-prev {
+ margin-right: auto;
+ width: fit-content;
+ float: none;
+}
+.prev-next-area a.right-next {
+ margin-left: auto;
+ width: fit-content;
+ float: none;
+}
+
+/* CSAT widgets */
+#csat-inputs {
+ display: flex;
+ flex-direction: row;
+ align-items: center;
+}
+
+.csat-hidden {
+ display: none !important;
+}
+
+#csat-feedback-label {
+ color: #000;
+ font-weight: 500;
+}
+
+.csat-button {
+ margin-left: 16px;
+ padding: 8px 16px 8px 16px;
+ border-radius: 4px;
+ border: 1px solid #D2DCE6;
+ background: #FFF;
+ display: flex;
+ flex-direction: row;
+ align-items: center;
+ justify-content: center;
+ cursor: pointer;
+ width: 85px;
+}
+
+#csat-textarea-group {
+ display: flex;
+ flex-direction: column;
+}
+
+#csat-submit {
+ margin-left: auto;
+ font-weight: 700;
+ border: none;
+ margin-top: 12px;
+ cursor: pointer;
+}
+
+#csat-feedback-received {
+ display: flex;
+ flex-direction: row;
+ align-items: center;
+ justify-content: center;
+}
+
+.csat-button-active {
+ border: 1px solid #000;
+}
+
+.csat-icon {
+ margin-right: 4px;
+}
+
+footer.col.footer {
+ display: flex;
+ flex-direction: row;
+}
+
+footer.col.footer > p {
+ margin-left: auto;
+}
+
+#csat {
+ min-width: 60%;
+}
+
+#csat-textarea {
+ resize: none;
+}
+
+
+/* Ray Assistant */
+
+.container-xl.blurred {
+ filter: blur(5px);
+}
+
+.chat-widget {
+ position: fixed;
+ bottom: 10px;
+ right: 10px;
+ z-index: 1000;
+}
+
+.chat-popup {
+ display: none;
+ position: fixed;
+ top: 20%;
+ left: 50%;
+ transform: translate(-50%, -20%);
+ width: 50%;
+ height: 70%;
+ background-color: white;
+ border: 1px solid #ccc;
+ border-radius: 10px;
+ box-shadow: 0 5px 10px rgba(0,0,0,0.1);
+ z-index: 1001;
+ max-height: 1000px;
+ overflow: hidden;
+ padding-bottom: 40px;
+}
+
+.chatFooter {
+ position: absolute;
+ bottom: 0;
+ right: 0;
+ width: 100%;
+ background-color: #f8f9fa;
+}
+
+#openChatBtn {
+ background-color: #000;
+ color: #fff;
+ width: 70px;
+ height: 70px;
+ border-radius: 10px;
+ border: none;
+ display: flex;
+ align-items: center;
+ justify-content: center;
+}
+
+#closeChatBtn {
+ border: none;
+ background-color: transparent;
+ color: #000;
+ font-size: 1.2em;
+}
+
+#closeChatBtn:hover {
+ color: #888;
+}
+
+.chatHeader {
+ display: flex;
+ justify-content: space-between;
+ align-items: center;
+}
+
+.chatContentContainer {
+ padding: 15px;
+ max-height: calc(100% - 80px);
+ overflow-y: auto;
+}
+
+.chatContentContainer input {
+ margin-top: 10px;
+ margin-bottom: 10px;
+}
+
+#result{
+ padding: 15px;
+ border-radius: 10px;
+ margin-top: 10px;
+ margin-bottom: 10px;
+ background-color: #f8f9fa;
+ max-height: calc(100% - 20px);
+ overflow-y: auto;
+}
+
+.chatContentContainer textarea {
+ flex-grow: 1;
+ min-width: 50px;
+ max-height: 40px;
+ resize: none;
+}
+
+.searchBtn {
+ white-space: nowrap;
+}
+
+.input-group {
+ display: flex;
+ align-items: stretch;
+}
+
+/* Kapa Ask AI button */
+#kapa-widget-container figure {
+ padding: 0 !important;
+ }
+
+ .mantine-Modal-root figure {
+ padding: 0 !important;
+ }
+
+@font-face {
+ font-family: "Linux Biolinum Keyboard";
+ src: url(../fonts/LinBiolinum_Kah.ttf);
+}
+
+.keys {
+ font-family: "Linux Biolinum Keyboard", sans-serif;
+}
+
+.bd-article-container h1, .bd-article-container h2, .bd-article-container h3, .bd-article-container h4, .bd-article-container h5, .bd-article-container p.caption {
+ color: black;
+}
diff --git a/docs/_static/css/examples.css b/docs/_static/css/examples.css
new file mode 100644
index 000000000..d1f416d15
--- /dev/null
+++ b/docs/_static/css/examples.css
@@ -0,0 +1,218 @@
+
+#site-navigation {
+ width: 330px !important;
+ border-right: none;
+ margin-left: 32px;
+ overflow-y: auto;
+ max-height: calc(100vh - var(--sidebar-top));
+ position: sticky;
+ top: var(--sidebar-top) !important;
+ z-index: 1000;
+}
+
+#site-navigation h5 {
+ font-size: 16px;
+ font-weight: 600;
+ color: #000;
+}
+
+#site-navigation h6 {
+ font-size: 14px;
+ font-weight: 600;
+ color: #000;
+ text-transform: uppercase;
+}
+
+/* Hide the default sidebar content */
+#site-navigation > div.bd-sidebar__content {
+ display: none;
+}
+#site-navigation > div.rtd-footer-container {
+ display: none;
+}
+
+.searchDiv {
+ margin-bottom: 2em;
+}
+
+#searchInput {
+ width: 100%;
+ color: #5F6469;
+ border: 1px solid #D2DCE6;
+ height: 50px;
+ border-radius: 4px;
+ background-color: #F9FAFB;
+ background-image: url("data:image/svg+xml,%3Csvg width='25' height='25' viewBox='0 0 25 25' fill='none' xmlns='http://www.w3.org/2000/svg'%3E%3Cg id='Systems / search-line' clip-path='url(%23clip0_1_150)'%3E%3Crect width='24' height='24' transform='translate(0.398529 0.0546875)' fill='%23F9FAFB'/%3E%3Cg id='Group'%3E%3Cpath id='Vector' d='M18.4295 16.6717L22.7125 20.9537L21.2975 22.3687L17.0155 18.0857C15.4223 19.3629 13.4405 20.0576 11.3985 20.0547C6.43053 20.0547 2.39853 16.0227 2.39853 11.0547C2.39853 6.08669 6.43053 2.05469 11.3985 2.05469C16.3665 2.05469 20.3985 6.08669 20.3985 11.0547C20.4014 13.0967 19.7068 15.0784 18.4295 16.6717ZM16.4235 15.9297C17.6926 14.6246 18.4014 12.8751 18.3985 11.0547C18.3985 7.18669 15.2655 4.05469 11.3985 4.05469C7.53053 4.05469 4.39853 7.18669 4.39853 11.0547C4.39853 14.9217 7.53053 18.0547 11.3985 18.0547C13.219 18.0576 14.9684 17.3488 16.2735 16.0797L16.4235 15.9297V15.9297Z' fill='%238C9196'/%3E%3C/g%3E%3C/g%3E%3Cdefs%3E%3CclipPath id='clip0_1_150'%3E%3Crect width='24' height='24' fill='white' transform='translate(0.398529 0.0546875)'/%3E%3C/clipPath%3E%3C/defs%3E%3C/svg%3E%0A");
+ background-repeat: no-repeat;
+ background-position-x: 0.5em;
+ background-position-y: center;
+ background-size: 1.5em;
+ padding-left: 3em;
+}
+
+#searchInput::placeholder {
+ color: #5F6469;
+ opacity: 1;
+}
+
+.tag {
+ margin-bottom: 5px;
+ font-size: small;
+ color: #000000;
+ border: 1px solid #D2DCE6;
+ border-radius: 14px;
+ display: flex;
+ flex-direction: row;
+ align-items: center;
+ width: fit-content;
+ gap: 1em;
+}
+
+.tag.btn-outline-primary {
+ color: #000000;
+ padding: 3px 12px 3px 12px;
+ line-height: 20px;
+}
+
+.tag-btn-wrapper {
+ display: flex;
+ flex-direction: row;
+ flex-wrap: wrap;
+ gap: 1em;
+}
+
+div.sd-container-fluid.docutils > div {
+ gap: var(--ray-example-gallery-gap-y) var(--ray-example-gallery-gap-x);
+ display: grid;
+ grid-template-columns: 1fr;
+}
+
+/* Reflow to a 2-column format for normal screens */
+@media screen and (min-width: 768px) {
+ div.sd-container-fluid.docutils > div {
+ grid-template-columns: 1fr 1fr;
+ }
+}
+
+div.gallery-item {
+ width: auto;
+}
+
+div.gallery-item > div.sd-card {
+ border-radius: 8px;
+ box-shadow: 0px 4px 10px 0px rgba(0, 0, 0, 0.05) !important;
+}
+
+/* Example gallery "Tutorial" title */
+div.sd-card-title > span.sd-bg-success.sd-bg-text-success {
+ color: #2F80ED !important;
+ font-weight: 500;
+ background: linear-gradient(180deg, rgba(25, 177, 226, 0.2) 0%, rgba(0, 109, 255, 0.2) 100%);
+ background-color: initial !important;
+}
+
+/* Example gallery "Code example" title */
+div.sd-card-title > span.sd-bg-secondary.sd-bg-text-secondary {
+ color: #219653 !important;
+ font-weight: 500;
+ background: linear-gradient(180deg, rgba(29, 151, 108, 0.2) 0%, rgba(0, 226, 147, 0.2) 100%);
+ background-color: initial !important;
+}
+
+/* Example gallery "Blog" title */
+div.sd-card-title > span.sd-bg-primary.sd-bg-text-primary {
+ color: #F2994A !important;
+ font-weight: 500;
+ background: linear-gradient(180deg, rgba(255, 230, 5, 0.2) 0%, rgba(255, 185, 80, 0.2) 100%);
+ background-color: initial !important;
+}
+
+/* Example gallery "Video" title */
+div.sd-card-title > span.sd-bg-warning.sd-bg-text-warning {
+ color: #EB5757 !important;
+ font-weight: 500;
+ background: linear-gradient(180deg, rgba(150, 7, 7, 0.2) 0%, rgba(255, 115, 115, 0.2) 100%);
+ background-color: initial !important;
+}
+
+/* Example gallery "Course" title */
+div.sd-card-title > span.sd-bg-info.sd-bg-text-info {
+ color: #7A64FF !important;
+ font-weight: 500;
+ background: linear-gradient(180deg, rgba(53, 25, 226, 0.2) 0%, rgba(183, 149, 255, 0.2) 100%);
+ background-color: initial !important;
+}
+
+div.sd-card-body > p.sd-card-text > a {
+ text-align: initial;
+}
+
+div.sd-card-body > p.sd-card-text > a > span {
+ color: rgb(81, 81, 81);
+}
+
+#main-content {
+ max-width: 100%;
+}
+
+#noMatches {
+ display: flex;
+ flex-direction: column;
+ align-items: center;
+ justify-content: center;
+}
+
+#noMatchesInnerContent {
+ display: flex;
+ flex-direction: column;
+ align-items: center;
+ justify-content: center;
+}
+
+#noMatches.hidden,.gallery-item.hidden {
+ display: none !important;
+}
+
+.btn-primary {
+ color: #004293;
+ background: rgba(61, 138, 233, 0.20);
+ padding: 3px 12px 3px 12px;
+ border: 1px solid #D2DCE6;
+}
+
+button.try-anyscale {
+ background-color: initial !important;
+ width: fit-content;
+ padding: 0 !important;
+ margin-left: auto !important;
+ float: initial !important;
+}
+
+button.try-anyscale > svg {
+ display: none;
+}
+
+button.try-anyscale > i {
+ display: none;
+}
+
+button.try-anyscale > span {
+ margin: 0;
+ text-decoration-line: underline;
+ font-weight: 500;
+ color: #000;
+}
+
+.top-nav-content {
+ justify-content: initial;
+}
+
+/* Hide nav bar that has github, fullscreen, and print icons */
+div.header-article.row.sticky-top.noprint {
+ display: none !important;
+}
+
+/* Hide the footer with 'prev article' and 'next article' buttons */
+.footer-article.hidden {
+ display: none !important;
+}
diff --git a/docs/_static/css/termynal.css b/docs/_static/css/termynal.css
new file mode 100644
index 000000000..391a48078
--- /dev/null
+++ b/docs/_static/css/termynal.css
@@ -0,0 +1,108 @@
+/**
+ * termynal.js
+ *
+ * @author Ines Montani
+
+The difference between **Public Prompts** and **Private Prompts** is that Prompts in **Public Prompts** space can be viewed and used by all users, while prompts in **Private Prompts** space can only be viewed and used by the owner.
+
+### Create Prompt
+
+Click the "Add Prompts" button to pop up the following subpage:
+
+
+
+**Scene**: It is assumed here that when we have a lot of Prompts, we often classify the Prompts according to scene, such as Prompts in the chat knowledge scene, Prompts in the chat data scene, Prompts in the chat normal scene, etc.
+
+**Sub Scene**: Continuing with the above, assuming that we have a lot of Prompts, scene classification alone is not enough. For example, in the chat data scenario, there can be many types of sub-scene: anomaly recognition sub scene, attribution analysis sub scene, etc. sub scene is used to distinguish subcategories under each scene.
+
+**Name**: Considering that a Prompt generally contains a lot of content, for ease of use and easy search, we need to name the Prompt. Note: The name of the Prompt is not allowed to be repeated. Name is the unique key that identifies a Prompt.
+
+**Content**: Here is the actual Prompt content that will be input to LLM.
+
+### Edit Prompt
+
+Existing Prompts can be edited. Note that except **name**, other items can be modified.
+
+
+
+### Delete Prompt
+
+Ordinary users can only delete Prompts created by themselves in the private Prompts space. Administrator users can delete Prompts in public Prompts spaces and private Prompts spaces.
+
+
+## Use Prompt
+
+Users can find and use Prompts next to the input boxes in each scene. Click to view all contents of Prompts library.
+
+✓ Hover the mouse over each Prompt to preview the Prompt content.
+✓ Click Prompt to automatically fill in the Prompt content in the input box.
+
+
+
\ No newline at end of file
diff --git a/docs/reference.md b/docs/reference.md
deleted file mode 100644
index 4a938e09d..000000000
--- a/docs/reference.md
+++ /dev/null
@@ -1 +0,0 @@
-# Reference
\ No newline at end of file
diff --git a/docs/requirements.txt b/docs/requirements.txt
index 8cee184af..d0a125868 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -9,9 +9,9 @@ sphinx_book_theme
sphinx_rtd_theme==1.0.0
sphinx-typlog-theme==0.8.0
sphinx-panels
+sphinx-tabs==3.4.0
toml
myst_nb
sphinx_copybutton
pydata-sphinx-theme==0.13.1
-pydantic-settings
furo
\ No newline at end of file
diff --git a/examples/app.py b/examples/app.py
deleted file mode 100644
index 07f8a5a51..000000000
--- a/examples/app.py
+++ /dev/null
@@ -1,74 +0,0 @@
-#!/usr/bin/env python3
-# -*- coding:utf-8 -*-
-
-import gradio as gr
-from langchain.agents import AgentType, initialize_agent, load_tools
-from llama_index import (
- Document,
- GPTVectorStoreIndex,
- LangchainEmbedding,
- LLMPredictor,
- ServiceContext,
-)
-
-from pilot.model.llm_out.vicuna_llm import VicunaEmbeddingLLM, VicunaRequestLLM
-
-
-def agent_demo():
- llm = VicunaRequestLLM()
-
- tools = load_tools(["python_repl"], llm=llm)
- agent = initialize_agent(
- tools, llm, agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION, verbose=True
- )
- agent.run("Write a SQL script that Query 'select count(1)!'")
-
-
-def knowledged_qa_demo(text_list):
- llm_predictor = LLMPredictor(llm=VicunaRequestLLM())
- hfemb = VicunaEmbeddingLLM()
- embed_model = LangchainEmbedding(hfemb)
- documents = [Document(t) for t in text_list]
-
- service_context = ServiceContext.from_defaults(
- llm_predictor=llm_predictor, embed_model=embed_model
- )
- index = GPTVectorStoreIndex.from_documents(
- documents, service_context=service_context
- )
- return index
-
-
-def get_answer(q):
- base_knowledge = """ """
- text_list = [base_knowledge]
- index = knowledged_qa_demo(text_list)
- response = index.query(q)
- return response.response
-
-
-def get_similar(q):
- from pilot.vector_store.extract_tovec import knownledge_tovec_st
-
- docsearch = knownledge_tovec_st("./datasets/plan.md")
- docs = docsearch.similarity_search_with_score(q, k=1)
-
- for doc in docs:
- dc, s = doc
- print(s)
- yield dc.page_content
-
-
-if __name__ == "__main__":
- # agent_demo()
-
- with gr.Blocks() as demo:
- gr.Markdown("数据库智能助手")
- with gr.Tab("知识问答"):
- text_input = gr.TextArea()
- text_output = gr.TextArea()
- text_button = gr.Button()
-
- text_button.click(get_similar, inputs=text_input, outputs=text_output)
-
- demo.queue(concurrency_count=3).launch(server_name="0.0.0.0")
diff --git a/examples/awel/simple_chat_dag_example.py b/examples/awel/simple_chat_dag_example.py
new file mode 100644
index 000000000..b53c2415e
--- /dev/null
+++ b/examples/awel/simple_chat_dag_example.py
@@ -0,0 +1,56 @@
+"""AWEL: Simple chat dag example
+
+ DB-GPT will automatically load and execute the current file after startup.
+
+ Example:
+
+ .. code-block:: shell
+
+ curl -X POST http://127.0.0.1:5000/api/v1/awel/trigger/examples/simple_chat \
+ -H "Content-Type: application/json" -d '{
+ "model": "proxyllm",
+ "user_input": "hello"
+ }'
+"""
+from typing import Dict
+from pydantic import BaseModel, Field
+
+from pilot.awel import DAG, HttpTrigger, MapOperator
+from pilot.scene.base_message import ModelMessage
+from pilot.model.base import ModelOutput
+from pilot.model.operator.model_operator import ModelOperator
+
+
+class TriggerReqBody(BaseModel):
+ model: str = Field(..., description="Model name")
+ user_input: str = Field(..., description="User input")
+
+
+class RequestHandleOperator(MapOperator[TriggerReqBody, Dict]):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ async def map(self, input_value: TriggerReqBody) -> Dict:
+ hist = []
+ hist.append(ModelMessage.build_human_message(input_value.user_input))
+ hist = list(h.dict() for h in hist)
+ params = {
+ "prompt": input_value.user_input,
+ "messages": hist,
+ "model": input_value.model,
+ "echo": False,
+ }
+ print(f"Receive input value: {input_value}")
+ return params
+
+
+with DAG("dbgpt_awel_simple_dag_example") as dag:
+ # Receive http request and trigger dag to run.
+ trigger = HttpTrigger(
+ "/examples/simple_chat", methods="POST", request_body=TriggerReqBody
+ )
+ request_handle_task = RequestHandleOperator()
+ model_task = ModelOperator()
+ # type(out) == ModelOutput
+ model_parse_task = MapOperator(lambda out: out.to_dict())
+ trigger >> request_handle_task >> model_task >> model_parse_task
diff --git a/examples/awel/simple_dag_example.py b/examples/awel/simple_dag_example.py
new file mode 100644
index 000000000..bfc9b45b5
--- /dev/null
+++ b/examples/awel/simple_dag_example.py
@@ -0,0 +1,34 @@
+"""AWEL: Simple dag example
+
+ DB-GPT will automatically load and execute the current file after startup.
+
+ Example:
+
+ .. code-block:: shell
+
+ curl -X GET http://127.0.0.1:5000/api/v1/awel/trigger/examples/hello\?name\=zhangsan
+
+"""
+from pydantic import BaseModel, Field
+
+from pilot.awel import DAG, HttpTrigger, MapOperator
+
+
+class TriggerReqBody(BaseModel):
+ name: str = Field(..., description="User name")
+ age: int = Field(18, description="User age")
+
+
+class RequestHandleOperator(MapOperator[TriggerReqBody, str]):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ async def map(self, input_value: TriggerReqBody) -> str:
+ print(f"Receive input value: {input_value}")
+ return f"Hello, {input_value.name}, your age is {input_value.age}"
+
+
+with DAG("simple_dag_example") as dag:
+ trigger = HttpTrigger("/examples/hello", request_body=TriggerReqBody)
+ map_node = RequestHandleOperator()
+ trigger >> map_node
diff --git a/examples/awel/simple_rag_example.py b/examples/awel/simple_rag_example.py
new file mode 100644
index 000000000..c7cd934dc
--- /dev/null
+++ b/examples/awel/simple_rag_example.py
@@ -0,0 +1,73 @@
+"""AWEL: Simple rag example
+
+ DB-GPT will automatically load and execute the current file after startup.
+
+ Example:
+
+ .. code-block:: shell
+
+ curl -X POST http://127.0.0.1:5000/api/v1/awel/trigger/examples/simple_rag \
+ -H "Content-Type: application/json" -d '{
+ "conv_uid": "36f0e992-8825-11ee-8638-0242ac150003",
+ "model_name": "proxyllm",
+ "chat_mode": "chat_knowledge",
+ "user_input": "What is DB-GPT?",
+ "select_param": "default"
+ }'
+
+"""
+
+from pilot.awel import HttpTrigger, DAG, MapOperator
+from pilot.scene.operator._experimental import (
+ ChatContext,
+ PromptManagerOperator,
+ ChatHistoryStorageOperator,
+ ChatHistoryOperator,
+ EmbeddingEngingOperator,
+ BaseChatOperator,
+)
+from pilot.scene.base import ChatScene
+from pilot.openapi.api_view_model import ConversationVo
+from pilot.model.base import ModelOutput
+from pilot.model.operator.model_operator import ModelOperator
+
+
+class RequestParseOperator(MapOperator[ConversationVo, ChatContext]):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ async def map(self, input_value: ConversationVo) -> ChatContext:
+ return ChatContext(
+ current_user_input=input_value.user_input,
+ model_name=input_value.model_name,
+ chat_session_id=input_value.conv_uid,
+ select_param=input_value.select_param,
+ chat_scene=ChatScene.ChatKnowledge,
+ )
+
+
+with DAG("simple_rag_example") as dag:
+ trigger_task = HttpTrigger(
+ "/examples/simple_rag", methods="POST", request_body=ConversationVo
+ )
+ req_parse_task = RequestParseOperator()
+ # TODO should register prompt template first
+ prompt_task = PromptManagerOperator()
+ history_storage_task = ChatHistoryStorageOperator()
+ history_task = ChatHistoryOperator()
+ embedding_task = EmbeddingEngingOperator()
+ chat_task = BaseChatOperator()
+ model_task = ModelOperator()
+ output_parser_task = MapOperator(lambda out: out.to_dict()["text"])
+
+ (
+ trigger_task
+ >> req_parse_task
+ >> prompt_task
+ >> history_storage_task
+ >> history_task
+ >> embedding_task
+ >> chat_task
+ >> model_task
+ >> output_parser_task
+ )
diff --git a/examples/embdserver.py b/examples/embdserver.py
deleted file mode 100644
index ae0dfcae8..000000000
--- a/examples/embdserver.py
+++ /dev/null
@@ -1,82 +0,0 @@
-#!/usr/bin/env python3
-# -*- coding:utf-8 -*-
-
-import json
-import os
-import sys
-from urllib.parse import urljoin
-
-import gradio as gr
-import requests
-
-ROOT_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-sys.path.append(ROOT_PATH)
-
-
-from langchain.prompts import PromptTemplate
-
-from pilot.configs.config import Config
-from pilot.conversation import conv_qa_prompt_template, conv_templates
-
-llmstream_stream_path = "generate_stream"
-
-CFG = Config()
-
-
-def generate(query):
- template_name = "conv_one_shot"
- state = conv_templates[template_name].copy()
-
- # pt = PromptTemplate(
- # template=conv_qa_prompt_template,
- # input_variables=["context", "question"]
- # )
-
- # result = pt.format(context="This page covers how to use the Chroma ecosystem within LangChain. It is broken into two parts: installation and setup, and then references to specific Chroma wrappers.",
- # question=query)
-
- # print(result)
-
- state.append_message(state.roles[0], query)
- state.append_message(state.roles[1], None)
-
- prompt = state.get_prompt()
- params = {
- "model": "chatglm-6b",
- "prompt": prompt,
- "temperature": 1.0,
- "max_new_tokens": 1024,
- "stop": "###",
- }
-
- response = requests.post(
- url=urljoin(CFG.MODEL_SERVER, llmstream_stream_path), data=json.dumps(params)
- )
-
- skip_echo_len = len(params["prompt"]) + 1 - params["prompt"].count("") * 3
-
- for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
- if chunk:
- data = json.loads(chunk.decode())
- if data["error_code"] == 0:
- if "vicuna" in CFG.LLM_MODEL:
- output = data["text"][skip_echo_len:].strip()
- else:
- output = data["text"].strip()
-
- state.messages[-1][-1] = output + "▌"
- yield (output)
-
-
-if __name__ == "__main__":
- print(CFG.LLM_MODEL)
- with gr.Blocks() as demo:
- gr.Markdown("数据库SQL生成助手")
- with gr.Tab("SQL生成"):
- text_input = gr.TextArea()
- text_output = gr.TextArea()
- text_button = gr.Button("提交")
-
- text_button.click(generate, inputs=text_input, outputs=text_output)
-
- demo.queue(concurrency_count=3).launch(server_name="0.0.0.0")
diff --git a/examples/gpt_index.py b/examples/gpt_index.py
deleted file mode 100644
index a4683af1a..000000000
--- a/examples/gpt_index.py
+++ /dev/null
@@ -1,19 +0,0 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-
-import logging
-import sys
-
-from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader
-
-logging.basicConfig(stream=sys.stdout, level=logging.INFO)
-logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
-
-# read the document of data dir
-documents = SimpleDirectoryReader("data").load_data()
-# split the document to chunk, max token size=500, convert chunk to vector
-
-index = GPTVectorStoreIndex(documents)
-
-# save index
-index.save_to_disk("index.json")
diff --git a/examples/gradio_test.py b/examples/gradio_test.py
deleted file mode 100644
index 593c6c1f4..000000000
--- a/examples/gradio_test.py
+++ /dev/null
@@ -1,21 +0,0 @@
-#!/usr/bin/env python3
-# -*- coding:utf-8 -*-
-
-import gradio as gr
-
-
-def change_tab():
- return gr.Tabs.update(selected=1)
-
-
-with gr.Blocks() as demo:
- with gr.Tabs() as tabs:
- with gr.TabItem("Train", id=0):
- t = gr.Textbox()
- with gr.TabItem("Inference", id=1):
- i = gr.Image()
-
- btn = gr.Button()
- btn.click(change_tab, None, tabs)
-
-demo.launch()
diff --git a/examples/knowledge_embedding/csv_embedding_test.py b/examples/knowledge_embedding/csv_embedding_test.py
deleted file mode 100644
index dcf4873b2..000000000
--- a/examples/knowledge_embedding/csv_embedding_test.py
+++ /dev/null
@@ -1,18 +0,0 @@
-from pilot.embedding_engine.csv_embedding import CSVEmbedding
-
-# path = "/Users/chenketing/Downloads/share_ireserve双写数据异常2.xlsx"
-path = "xx.csv"
-model_name = "your_path/all-MiniLM-L6-v2"
-vector_store_path = "your_path/"
-
-
-pdf_embedding = CSVEmbedding(
- file_path=path,
- model_name=model_name,
- vector_store_config={
- "vector_store_name": "url",
- "vector_store_path": "vector_store_path",
- },
-)
-pdf_embedding.source_embedding()
-print("success")
diff --git a/examples/knowledge_embedding/pdf_embedding_test.py b/examples/knowledge_embedding/pdf_embedding_test.py
deleted file mode 100644
index ef0e1d87e..000000000
--- a/examples/knowledge_embedding/pdf_embedding_test.py
+++ /dev/null
@@ -1,18 +0,0 @@
-from pilot.embedding_engine.pdf_embedding import PDFEmbedding
-
-path = "xxx.pdf"
-path = "your_path/OceanBase-数据库-V4.1.0-应用开发.pdf"
-model_name = "your_path/all-MiniLM-L6-v2"
-vector_store_path = "your_path/"
-
-
-pdf_embedding = PDFEmbedding(
- file_path=path,
- model_name=model_name,
- vector_store_config={
- "vector_store_name": "ob-pdf",
- "vector_store_path": vector_store_path,
- },
-)
-pdf_embedding.source_embedding()
-print("success")
diff --git a/examples/knowledge_embedding/url_embedding_test.py b/examples/knowledge_embedding/url_embedding_test.py
deleted file mode 100644
index c702fd1f7..000000000
--- a/examples/knowledge_embedding/url_embedding_test.py
+++ /dev/null
@@ -1,17 +0,0 @@
-from pilot.embedding_engine.url_embedding import URLEmbedding
-
-path = "https://www.understandingwar.org/backgrounder/russian-offensive-campaign-assessment-february-8-2023"
-model_name = "your_path/all-MiniLM-L6-v2"
-vector_store_path = "your_path"
-
-
-pdf_embedding = URLEmbedding(
- file_path=path,
- model_name=model_name,
- vector_store_config={
- "vector_store_name": "url",
- "vector_store_path": "vector_store_path",
- },
-)
-pdf_embedding.source_embedding()
-print("success")
diff --git a/examples/proxy_example.py b/examples/proxy_example.py
deleted file mode 100644
index a3d2f3bc4..000000000
--- a/examples/proxy_example.py
+++ /dev/null
@@ -1,67 +0,0 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-
-import dashscope
-import requests
-import hashlib
-from http import HTTPStatus
-from dashscope import Generation
-
-
-def call_with_messages():
- messages = [
- {"role": "system", "content": "你是生活助手机器人。"},
- {"role": "user", "content": "如何做西红柿鸡蛋?"},
- ]
- gen = Generation()
- response = gen.call(
- Generation.Models.qwen_turbo,
- messages=messages,
- stream=True,
- top_p=0.8,
- result_format="message", # set the result to be "message" format.
- )
-
- for response in response:
- # The response status_code is HTTPStatus.OK indicate success,
- # otherwise indicate request is failed, you can get error code
- # and message from code and message.
- if response.status_code == HTTPStatus.OK:
- print(response.output) # The output text
- print(response.usage) # The usage information
- else:
- print(response.code) # The error code.
- print(response.message) # The error message.
-
-
-def build_access_token(api_key: str, secret_key: str) -> str:
- """
- Generate Access token according AK, SK
- """
-
- url = "https://aip.baidubce.com/oauth/2.0/token"
- params = {
- "grant_type": "client_credentials",
- "client_id": api_key,
- "client_secret": secret_key,
- }
-
- res = requests.get(url=url, params=params)
-
- if res.status_code == 200:
- return res.json().get("access_token")
-
-
-def _calculate_md5(text: str) -> str:
- md5 = hashlib.md5()
- md5.update(text.encode("utf-8"))
- encrypted = md5.hexdigest()
- return encrypted
-
-
-def baichuan_call():
- url = "https://api.baichuan-ai.com/v1/stream/chat"
-
-
-if __name__ == "__main__":
- call_with_messages()
diff --git a/examples/t5_example.py b/examples/t5_example.py
deleted file mode 100644
index b49e79d4e..000000000
--- a/examples/t5_example.py
+++ /dev/null
@@ -1,257 +0,0 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-
-import torch
-from langchain.embeddings.huggingface import HuggingFaceEmbeddings
-from langchain.llms.base import LLM
-from llama_index import (
- GPTListIndex,
- GPTVectorStoreIndex,
- LangchainEmbedding,
- LLMPredictor,
- PromptHelper,
- SimpleDirectoryReader,
-)
-from transformers import pipeline
-
-
-class FlanLLM(LLM):
- model_name = "google/flan-t5-large"
- pipeline = pipeline(
- "text2text-generation",
- model=model_name,
- device=0,
- model_kwargs={"torch_dtype": torch.bfloat16},
- )
-
- def _call(self, prompt, stop=None):
- return self.pipeline(prompt, max_length=9999)[0]["generated_text"]
-
- def _identifying_params(self):
- return {"name_of_model": self.model_name}
-
- def _llm_type(self):
- return "custome"
-
-
-llm_predictor = LLMPredictor(llm=FlanLLM())
-hfemb = HuggingFaceEmbeddings()
-embed_model = LangchainEmbedding(hfemb)
-
-text1 = """
- 执行计划是对一条 SQL 查询语句在数据库中执行过程的描述。用户可以通过 EXPLAIN 命令查看优化器针对指定 SQL 生成的逻辑执行计划。
-
-如果要分析某条 SQL 的性能问题,通常需要先查看 SQL 的执行计划,排查每一步 SQL 执行是否存在问题。所以读懂执行计划是 SQL 优化的先决条件,而了解执行计划的算子是理解 EXPLAIN 命令的关键。
-
-OceanBase 数据库的执行计划命令有三种模式:EXPLAIN BASIC、EXPLAIN 和 EXPLAIN EXTENDED。这三种模式对执行计划展现不同粒度的细节信息:
-
-EXPLAIN BASIC 命令用于最基本的计划展示。
-
-EXPLAIN EXTENDED 命令用于最详细的计划展示(通常在排查问题时使用这种展示模式)。
-
-EXPLAIN 命令所展示的信息可以帮助普通用户了解整个计划的执行方式。
-
-EXPLAIN 命令格式如下:
-EXPLAIN [BASIC | EXTENDED | PARTITIONS | FORMAT = format_name] [PRETTY | PRETTY_COLOR] explainable_stmt
-format_name:
- { TRADITIONAL | JSON }
-explainable_stmt:
- { SELECT statement
- | DELETE statement
- | INSERT statement
- | REPLACE statement
- | UPDATE statement }
-
-
-EXPLAIN 命令适用于 SELECT、DELETE、INSERT、REPLACE 和 UPDATE 语句,显示优化器所提供的有关语句执行计划的信息,包括如何处理该语句,如何联接表以及以何种顺序联接表等信息。
-
-一般来说,可以使用 EXPLAIN EXTENDED 命令,将表扫描的范围段展示出来。使用 EXPLAIN OUTLINE 命令可以显示 Outline 信息。
-
-FORMAT 选项可用于选择输出格式。TRADITIONAL 表示以表格格式显示输出,这也是默认设置。JSON 表示以 JSON 格式显示信息。
-
-使用 EXPLAIN PARTITITIONS 也可用于检查涉及分区表的查询。如果检查针对非分区表的查询,则不会产生错误,但 PARTIONS 列的值始终为 NULL。
-
-对于复杂的执行计划,可以使用 PRETTY 或者 PRETTY_COLOR 选项将计划树中的父节点和子节点使用树线或彩色树线连接起来,使得执行计划展示更方便阅读。示例如下:
-obclient> CREATE TABLE p1table(c1 INT ,c2 INT) PARTITION BY HASH(c1) PARTITIONS 2;
-Query OK, 0 rows affected
-
-obclient> CREATE TABLE p2table(c1 INT ,c2 INT) PARTITION BY HASH(c1) PARTITIONS 4;
-Query OK, 0 rows affected
-
-obclient> EXPLAIN EXTENDED PRETTY_COLOR SELECT * FROM p1table p1 JOIN p2table p2 ON p1.c1=p2.c2\G
-*************************** 1. row ***************************
-Query Plan: ==========================================================
-|ID|OPERATOR |NAME |EST. ROWS|COST|
-----------------------------------------------------------
-|0 |PX COORDINATOR | |1 |278 |
-|1 | EXCHANGE OUT DISTR |:EX10001|1 |277 |
-|2 | HASH JOIN | |1 |276 |
-|3 | ├PX PARTITION ITERATOR | |1 |92 |
-|4 | │ TABLE SCAN |P1 |1 |92 |
-|5 | └EXCHANGE IN DISTR | |1 |184 |
-|6 | EXCHANGE OUT DISTR (PKEY)|:EX10000|1 |184 |
-|7 | PX PARTITION ITERATOR | |1 |183 |
-|8 | TABLE SCAN |P2 |1 |183 |
-==========================================================
-
-Outputs & filters:
--------------------------------------
- 0 - output([INTERNAL_FUNCTION(P1.C1, P1.C2, P2.C1, P2.C2)]), filter(nil)
- 1 - output([INTERNAL_FUNCTION(P1.C1, P1.C2, P2.C1, P2.C2)]), filter(nil), dop=1
- 2 - output([P1.C1], [P2.C2], [P1.C2], [P2.C1]), filter(nil),
- equal_conds([P1.C1 = P2.C2]), other_conds(nil)
- 3 - output([P1.C1], [P1.C2]), filter(nil)
- 4 - output([P1.C1], [P1.C2]), filter(nil),
- access([P1.C1], [P1.C2]), partitions(p[0-1])
- 5 - output([P2.C2], [P2.C1]), filter(nil)
- 6 - (#keys=1, [P2.C2]), output([P2.C2], [P2.C1]), filter(nil), dop=1
- 7 - output([P2.C1], [P2.C2]), filter(nil)
- 8 - output([P2.C1], [P2.C2]), filter(nil),
- access([P2.C1], [P2.C2]), partitions(p[0-3])
-
-1 row in set
-
-
-
-
-## 执行计划形状与算子信息
-
-在数据库系统中,执行计划在内部通常是以树的形式来表示的,但是不同的数据库会选择不同的方式展示给用户。
-
-如下示例分别为 PostgreSQL 数据库、Oracle 数据库和 OceanBase 数据库对于 TPCDS Q3 的计划展示。
-
-```sql
-obclient> SELECT /*TPC-DS Q3*/ *
- FROM (SELECT dt.d_year,
- item.i_brand_id brand_id,
- item.i_brand brand,
- Sum(ss_net_profit) sum_agg
- FROM date_dim dt,
- store_sales,
- item
- WHERE dt.d_date_sk = store_sales.ss_sold_date_sk
- AND store_sales.ss_item_sk = item.i_item_sk
- AND item.i_manufact_id = 914
- AND dt.d_moy = 11
- GROUP BY dt.d_year,
- item.i_brand,
- item.i_brand_id
- ORDER BY dt.d_year,
- sum_agg DESC,
- brand_id)
- WHERE ROWNUM <= 100;
-
-PostgreSQL 数据库执行计划展示如下:
-Limit (cost=13986.86..13987.20 rows=27 width=91)
- Sort (cost=13986.86..13986.93 rows=27 width=65)
- Sort Key: dt.d_year, (sum(store_sales.ss_net_profit)), item.i_brand_id
- HashAggregate (cost=13985.95..13986.22 rows=27 width=65)
- Merge Join (cost=13884.21..13983.91 rows=204 width=65)
- Merge Cond: (dt.d_date_sk = store_sales.ss_sold_date_sk)
- Index Scan using date_dim_pkey on date_dim dt (cost=0.00..3494.62 rows=6080 width=8)
- Filter: (d_moy = 11)
- Sort (cost=12170.87..12177.27 rows=2560 width=65)
- Sort Key: store_sales.ss_sold_date_sk
- Nested Loop (cost=6.02..12025.94 rows=2560 width=65)
- Seq Scan on item (cost=0.00..1455.00 rows=16 width=59)
- Filter: (i_manufact_id = 914)
- Bitmap Heap Scan on store_sales (cost=6.02..658.94 rows=174 width=14)
- Recheck Cond: (ss_item_sk = item.i_item_sk)
- Bitmap Index Scan on store_sales_pkey (cost=0.00..5.97 rows=174 width=0)
- Index Cond: (ss_item_sk = item.i_item_sk)
-
-
-
-Oracle 数据库执行计划展示如下:
-Plan hash value: 2331821367
---------------------------------------------------------------------------------------------------
-| Id | Operation | Name | Rows | Bytes | Cost (%CPU)| Time |
---------------------------------------------------------------------------------------------------
-| 0 | SELECT STATEMENT | | 100 | 9100 | 3688 (1)| 00:00:01 |
-|* 1 | COUNT STOPKEY | | | | | |
-| 2 | VIEW | | 2736 | 243K| 3688 (1)| 00:00:01 |
-|* 3 | SORT ORDER BY STOPKEY | | 2736 | 256K| 3688 (1)| 00:00:01 |
-| 4 | HASH GROUP BY | | 2736 | 256K| 3688 (1)| 00:00:01 |
-|* 5 | HASH JOIN | | 2736 | 256K| 3686 (1)| 00:00:01 |
-|* 6 | TABLE ACCESS FULL | DATE_DIM | 6087 | 79131 | 376 (1)| 00:00:01 |
-| 7 | NESTED LOOPS | | 2865 | 232K| 3310 (1)| 00:00:01 |
-| 8 | NESTED LOOPS | | 2865 | 232K| 3310 (1)| 00:00:01 |
-|* 9 | TABLE ACCESS FULL | ITEM | 18 | 1188 | 375 (0)| 00:00:01 |
-|* 10 | INDEX RANGE SCAN | SYS_C0010069 | 159 | | 2 (0)| 00:00:01 |
-| 11 | TABLE ACCESS BY INDEX ROWID| STORE_SALES | 159 | 2703 | 163 (0)| 00:00:01 |
---------------------------------------------------------------------------------------------------
-
-OceanBase 数据库执行计划展示如下:
-|ID|OPERATOR |NAME |EST. ROWS|COST |
--------------------------------------------------------
-|0 |LIMIT | |100 |81141|
-|1 | TOP-N SORT | |100 |81127|
-|2 | HASH GROUP BY | |2924 |68551|
-|3 | HASH JOIN | |2924 |65004|
-|4 | SUBPLAN SCAN |VIEW1 |2953 |19070|
-|5 | HASH GROUP BY | |2953 |18662|
-|6 | NESTED-LOOP JOIN| |2953 |15080|
-|7 | TABLE SCAN |ITEM |19 |11841|
-|8 | TABLE SCAN |STORE_SALES|161 |73 |
-|9 | TABLE SCAN |DT |6088 |29401|
-=======================================================
-
-由示例可见,OceanBase 数据库的计划展示与 Oracle 数据库类似。
-
-OceanBase 数据库执行计划中的各列的含义如下:
-列名 含义
-ID 执行树按照前序遍历的方式得到的编号(从 0 开始)。
-OPERATOR 操作算子的名称。
-NAME 对应表操作的表名(索引名)。
-EST. ROWS 估算该操作算子的输出行数。
-COST 该操作算子的执行代价(微秒)。
-
-
-OceanBase 数据库 EXPLAIN 命令输出的第一部分是执行计划的树形结构展示。其中每一个操作在树中的层次通过其在 operator 中的缩进予以展示,层次最深的优先执行,层次相同的以特定算子的执行顺序为标准来执行。
-
-问题: update a not exists (b…)
-我一开始以为 B是驱动表,B的数据挺多的 后来看到NLAJ,是说左边的表关联右边的表
-所以这个的驱动表是不是实际是A,用A的匹配B的,这个理解有问题吗
-
-回答: 没错 A 驱动 B的
-
-问题: 光知道最下最右的是驱动表了 所以一开始搞得有点懵 :sweat_smile:
-
-回答: nlj应该原理应该都是左表(驱动表)的记录探测右表(被驱动表), 选哪张成为左表或右表就基于一些其他考量了,比如数据量, 而anti join/semi join只是对 not exist/exist的一种优化,相关的原理和资料网上可以查阅一下
-
-问题: 也就是nlj 就是按照之前理解的谁先执行 谁就是驱动表 也就是执行计划中的最右的表
-而anti join/semi join,谁在not exist左面,谁就是驱动表。这么理解对吧
-
-回答: nlj也是左表的表是驱动表,这个要了解下计划执行方面的基本原理,取左表的一行数据,再遍历右表,一旦满足连接条件,就可以返回数据
-anti/semi只是因为not exists/exist的语义只是返回左表数据,改成anti join是一种计划优化,连接的方式比子查询更优
-"""
-
-from llama_index import Document
-
-text_list = [text1]
-documents = [Document(t) for t in text_list]
-
-num_output = 250
-max_input_size = 512
-
-max_chunk_overlap = 20
-prompt_helper = PromptHelper(max_input_size, num_output, max_chunk_overlap)
-
-index = GPTListIndex(
- documents,
- embed_model=embed_model,
- llm_predictor=llm_predictor,
- prompt_helper=prompt_helper,
-)
-index.save_to_disk("index.json")
-
-
-if __name__ == "__main__":
- import logging
-
- logging.getLogger().setLevel(logging.CRITICAL)
- for d in documents:
- print(d)
-
- response = index.query("数据库的执行计划命令有多少?")
- print(response)
diff --git a/pilot/awel/__init__.py b/pilot/awel/__init__.py
new file mode 100644
index 000000000..3cfc3c2bc
--- /dev/null
+++ b/pilot/awel/__init__.py
@@ -0,0 +1,87 @@
+"""Agentic Workflow Expression Language (AWEL)
+
+Note:
+
+AWEL is still an experimental feature and only opens the lowest level API.
+The stability of this API cannot be guaranteed at present.
+
+"""
+
+from pilot.component import SystemApp
+
+from .dag.base import DAGContext, DAG
+
+from .operator.base import BaseOperator, WorkflowRunner
+from .operator.common_operator import (
+ JoinOperator,
+ ReduceStreamOperator,
+ MapOperator,
+ BranchOperator,
+ InputOperator,
+ BranchFunc,
+)
+
+from .operator.stream_operator import (
+ StreamifyAbsOperator,
+ UnstreamifyAbsOperator,
+ TransformStreamAbsOperator,
+)
+
+from .task.base import TaskState, TaskOutput, TaskContext, InputContext, InputSource
+from .task.task_impl import (
+ SimpleInputSource,
+ SimpleCallDataInputSource,
+ DefaultTaskContext,
+ DefaultInputContext,
+ SimpleTaskOutput,
+ SimpleStreamTaskOutput,
+ _is_async_iterator,
+)
+from .trigger.http_trigger import HttpTrigger
+from .runner.local_runner import DefaultWorkflowRunner
+
+__all__ = [
+ "initialize_awel",
+ "DAGContext",
+ "DAG",
+ "BaseOperator",
+ "JoinOperator",
+ "ReduceStreamOperator",
+ "MapOperator",
+ "BranchOperator",
+ "InputOperator",
+ "BranchFunc",
+ "WorkflowRunner",
+ "TaskState",
+ "TaskOutput",
+ "TaskContext",
+ "InputContext",
+ "InputSource",
+ "DefaultWorkflowRunner",
+ "SimpleInputSource",
+ "SimpleCallDataInputSource",
+ "DefaultTaskContext",
+ "DefaultInputContext",
+ "SimpleTaskOutput",
+ "SimpleStreamTaskOutput",
+ "StreamifyAbsOperator",
+ "UnstreamifyAbsOperator",
+ "TransformStreamAbsOperator",
+ "HttpTrigger",
+]
+
+
+def initialize_awel(system_app: SystemApp, dag_filepath: str):
+ from .dag.dag_manager import DAGManager
+ from .dag.base import DAGVar
+ from .trigger.trigger_manager import DefaultTriggerManager
+ from .operator.base import initialize_runner
+
+ DAGVar.set_current_system_app(system_app)
+
+ system_app.register(DefaultTriggerManager)
+ dag_manager = DAGManager(system_app, dag_filepath)
+ system_app.register_instance(dag_manager)
+ initialize_runner(DefaultWorkflowRunner())
+ # Load all dags
+ dag_manager.load_dags()
diff --git a/pilot/awel/base.py b/pilot/awel/base.py
new file mode 100644
index 000000000..97cb8ad05
--- /dev/null
+++ b/pilot/awel/base.py
@@ -0,0 +1,7 @@
+from abc import ABC, abstractmethod
+
+
+class Trigger(ABC):
+ @abstractmethod
+ async def trigger(self) -> None:
+ """Trigger the workflow or a specific operation in the workflow."""
diff --git a/pilot/server/componet_configs.py b/pilot/awel/dag/__init__.py
similarity index 100%
rename from pilot/server/componet_configs.py
rename to pilot/awel/dag/__init__.py
diff --git a/pilot/awel/dag/base.py b/pilot/awel/dag/base.py
new file mode 100644
index 000000000..ceb13c8ad
--- /dev/null
+++ b/pilot/awel/dag/base.py
@@ -0,0 +1,364 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Dict, List, Sequence, Union, Any, Set
+import uuid
+import contextvars
+import threading
+import asyncio
+import logging
+from collections import deque
+from functools import cache
+from concurrent.futures import Executor
+
+from pilot.component import SystemApp
+from ..resource.base import ResourceGroup
+from ..task.base import TaskContext
+
+logger = logging.getLogger(__name__)
+
+DependencyType = Union["DependencyMixin", Sequence["DependencyMixin"]]
+
+
+def _is_async_context():
+ try:
+ loop = asyncio.get_running_loop()
+ return asyncio.current_task(loop=loop) is not None
+ except RuntimeError:
+ return False
+
+
+class DependencyMixin(ABC):
+ @abstractmethod
+ def set_upstream(self, nodes: DependencyType) -> "DependencyMixin":
+ """Set one or more upstream nodes for this node.
+
+ Args:
+ nodes (DependencyType): Upstream nodes to be set to current node.
+
+ Returns:
+ DependencyMixin: Returns self to allow method chaining.
+
+ Raises:
+ ValueError: If no upstream nodes are provided or if an argument is not a DependencyMixin.
+ """
+
+ @abstractmethod
+ def set_downstream(self, nodes: DependencyType) -> "DependencyMixin":
+ """Set one or more downstream nodes for this node.
+
+ Args:
+ nodes (DependencyType): Downstream nodes to be set to current node.
+
+ Returns:
+ DependencyMixin: Returns self to allow method chaining.
+
+ Raises:
+ ValueError: If no downstream nodes are provided or if an argument is not a DependencyMixin.
+ """
+
+ def __lshift__(self, nodes: DependencyType) -> DependencyType:
+ """Implements self << nodes
+
+ Example:
+
+ .. code-block:: python
+
+ # means node.set_upstream(input_node)
+ node << input_node
+
+ # means node2.set_upstream([input_node])
+ node2 << [input_node]
+ """
+ self.set_upstream(nodes)
+ return nodes
+
+ def __rshift__(self, nodes: DependencyType) -> DependencyType:
+ """Implements self >> nodes
+
+ Example:
+
+ .. code-block:: python
+
+ # means node.set_downstream(next_node)
+ node >> next_node
+
+ # means node2.set_downstream([next_node])
+ node2 >> [next_node]
+
+ """
+ self.set_downstream(nodes)
+ return nodes
+
+ def __rrshift__(self, nodes: DependencyType) -> "DependencyMixin":
+ """Implements [node] >> self"""
+ self.__lshift__(nodes)
+ return self
+
+ def __rlshift__(self, nodes: DependencyType) -> "DependencyMixin":
+ """Implements [node] << self"""
+ self.__rshift__(nodes)
+ return self
+
+
+class DAGVar:
+ _thread_local = threading.local()
+ _async_local = contextvars.ContextVar("current_dag_stack", default=deque())
+ _system_app: SystemApp = None
+ _executor: Executor = None
+
+ @classmethod
+ def enter_dag(cls, dag) -> None:
+ is_async = _is_async_context()
+ if is_async:
+ stack = cls._async_local.get()
+ stack.append(dag)
+ cls._async_local.set(stack)
+ else:
+ if not hasattr(cls._thread_local, "current_dag_stack"):
+ cls._thread_local.current_dag_stack = deque()
+ cls._thread_local.current_dag_stack.append(dag)
+
+ @classmethod
+ def exit_dag(cls) -> None:
+ is_async = _is_async_context()
+ if is_async:
+ stack = cls._async_local.get()
+ if stack:
+ stack.pop()
+ cls._async_local.set(stack)
+ else:
+ if (
+ hasattr(cls._thread_local, "current_dag_stack")
+ and cls._thread_local.current_dag_stack
+ ):
+ cls._thread_local.current_dag_stack.pop()
+
+ @classmethod
+ def get_current_dag(cls) -> Optional["DAG"]:
+ is_async = _is_async_context()
+ if is_async:
+ stack = cls._async_local.get()
+ return stack[-1] if stack else None
+ else:
+ if (
+ hasattr(cls._thread_local, "current_dag_stack")
+ and cls._thread_local.current_dag_stack
+ ):
+ return cls._thread_local.current_dag_stack[-1]
+ return None
+
+ @classmethod
+ def get_current_system_app(cls) -> SystemApp:
+ if not cls._system_app:
+ raise RuntimeError("System APP not set for DAGVar")
+ return cls._system_app
+
+ @classmethod
+ def set_current_system_app(cls, system_app: SystemApp) -> None:
+ if cls._system_app:
+ logger.warn("System APP has already set, nothing to do")
+ else:
+ cls._system_app = system_app
+
+ @classmethod
+ def get_executor(cls) -> Executor:
+ return cls._executor
+
+ @classmethod
+ def set_executor(cls, executor: Executor) -> None:
+ cls._executor = executor
+
+
+class DAGNode(DependencyMixin, ABC):
+ resource_group: Optional[ResourceGroup] = None
+ """The resource group of current DAGNode"""
+
+ def __init__(
+ self,
+ dag: Optional["DAG"] = None,
+ node_id: Optional[str] = None,
+ node_name: Optional[str] = None,
+ system_app: Optional[SystemApp] = None,
+ executor: Optional[Executor] = None,
+ ) -> None:
+ super().__init__()
+ self._upstream: List["DAGNode"] = []
+ self._downstream: List["DAGNode"] = []
+ self._dag: Optional["DAG"] = dag or DAGVar.get_current_dag()
+ self._system_app: Optional[SystemApp] = (
+ system_app or DAGVar.get_current_system_app()
+ )
+ self._executor: Optional[Executor] = executor or DAGVar.get_executor()
+ if not node_id and self._dag:
+ node_id = self._dag._new_node_id()
+ self._node_id: str = node_id
+ self._node_name: str = node_name
+
+ @property
+ def node_id(self) -> str:
+ return self._node_id
+
+ @property
+ def system_app(self) -> SystemApp:
+ return self._system_app
+
+ def set_node_id(self, node_id: str) -> None:
+ self._node_id = node_id
+
+ def __hash__(self) -> int:
+ if self.node_id:
+ return hash(self.node_id)
+ else:
+ return super().__hash__()
+
+ def __eq__(self, other: Any) -> bool:
+ if not isinstance(other, DAGNode):
+ return False
+ return self.node_id == other.node_id
+
+ @property
+ def node_name(self) -> str:
+ return self._node_name
+
+ @property
+ def dag(self) -> "DAG":
+ return self._dag
+
+ def set_upstream(self, nodes: DependencyType) -> "DAGNode":
+ self.set_dependency(nodes)
+
+ def set_downstream(self, nodes: DependencyType) -> "DAGNode":
+ self.set_dependency(nodes, is_upstream=False)
+
+ @property
+ def upstream(self) -> List["DAGNode"]:
+ return self._upstream
+
+ @property
+ def downstream(self) -> List["DAGNode"]:
+ return self._downstream
+
+ def set_dependency(self, nodes: DependencyType, is_upstream: bool = True) -> None:
+ if not isinstance(nodes, Sequence):
+ nodes = [nodes]
+ if not all(isinstance(node, DAGNode) for node in nodes):
+ raise ValueError(
+ "all nodes to set dependency to current node must be instance of 'DAGNode'"
+ )
+ nodes: Sequence[DAGNode] = nodes
+ dags = set([node.dag for node in nodes if node.dag])
+ if self.dag:
+ dags.add(self.dag)
+ if not dags:
+ raise ValueError("set dependency to current node must in a DAG context")
+ if len(dags) != 1:
+ raise ValueError(
+ "set dependency to current node just support in one DAG context"
+ )
+ dag = dags.pop()
+ self._dag = dag
+
+ dag._append_node(self)
+ for node in nodes:
+ if is_upstream and node not in self.upstream:
+ node._dag = dag
+ dag._append_node(node)
+
+ self._upstream.append(node)
+ node._downstream.append(self)
+ elif node not in self._downstream:
+ node._dag = dag
+ dag._append_node(node)
+
+ self._downstream.append(node)
+ node._upstream.append(self)
+
+
+class DAGContext:
+ def __init__(self) -> None:
+ self._curr_task_ctx = None
+ self._share_data: Dict[str, Any] = {}
+
+ @property
+ def current_task_context(self) -> TaskContext:
+ return self._curr_task_ctx
+
+ def set_current_task_context(self, _curr_task_ctx: TaskContext) -> None:
+ self._curr_task_ctx = _curr_task_ctx
+
+ async def get_share_data(self, key: str) -> Any:
+ return self._share_data.get(key)
+
+ async def save_to_share_data(self, key: str, data: Any) -> None:
+ self._share_data[key] = data
+
+
+class DAG:
+ def __init__(
+ self, dag_id: str, resource_group: Optional[ResourceGroup] = None
+ ) -> None:
+ self._dag_id = dag_id
+ self.node_map: Dict[str, DAGNode] = {}
+ self._root_nodes: Set[DAGNode] = None
+ self._leaf_nodes: Set[DAGNode] = None
+ self._trigger_nodes: Set[DAGNode] = None
+
+ def _append_node(self, node: DAGNode) -> None:
+ self.node_map[node.node_id] = node
+ # clear cached nodes
+ self._root_nodes = None
+ self._leaf_nodes = None
+
+ def _new_node_id(self) -> str:
+ return str(uuid.uuid4())
+
+ @property
+ def dag_id(self) -> str:
+ return self._dag_id
+
+ def _build(self) -> None:
+ from ..operator.common_operator import TriggerOperator
+
+ nodes = set()
+ for _, node in self.node_map.items():
+ nodes = nodes.union(_get_nodes(node))
+ self._root_nodes = list(set(filter(lambda x: not x.upstream, nodes)))
+ self._leaf_nodes = list(set(filter(lambda x: not x.downstream, nodes)))
+ self._trigger_nodes = list(
+ set(filter(lambda x: isinstance(x, TriggerOperator), nodes))
+ )
+
+ @property
+ def root_nodes(self) -> List[DAGNode]:
+ if not self._root_nodes:
+ self._build()
+ return self._root_nodes
+
+ @property
+ def leaf_nodes(self) -> List[DAGNode]:
+ if not self._leaf_nodes:
+ self._build()
+ return self._leaf_nodes
+
+ @property
+ def trigger_nodes(self):
+ if not self._trigger_nodes:
+ self._build()
+ return self._trigger_nodes
+
+ def __enter__(self):
+ DAGVar.enter_dag(self)
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ DAGVar.exit_dag()
+
+
+def _get_nodes(node: DAGNode, is_upstream: Optional[bool] = True) -> set[DAGNode]:
+ nodes = set()
+ if not node:
+ return nodes
+ nodes.add(node)
+ stream_nodes = node.upstream if is_upstream else node.downstream
+ for node in stream_nodes:
+ nodes = nodes.union(_get_nodes(node, is_upstream))
+ return nodes
diff --git a/pilot/awel/dag/dag_manager.py b/pilot/awel/dag/dag_manager.py
new file mode 100644
index 000000000..58830e121
--- /dev/null
+++ b/pilot/awel/dag/dag_manager.py
@@ -0,0 +1,42 @@
+from typing import Dict, Optional
+import logging
+from pilot.component import BaseComponent, ComponentType, SystemApp
+from .loader import DAGLoader, LocalFileDAGLoader
+from .base import DAG
+
+logger = logging.getLogger(__name__)
+
+
+class DAGManager(BaseComponent):
+ name = ComponentType.AWEL_DAG_MANAGER
+
+ def __init__(self, system_app: SystemApp, dag_filepath: str):
+ super().__init__(system_app)
+ self.dag_loader = LocalFileDAGLoader(dag_filepath)
+ self.system_app = system_app
+ self.dag_map: Dict[str, DAG] = {}
+
+ def init_app(self, system_app: SystemApp):
+ self.system_app = system_app
+
+ def load_dags(self):
+ dags = self.dag_loader.load_dags()
+ triggers = []
+ for dag in dags:
+ dag_id = dag.dag_id
+ if dag_id in self.dag_map:
+ raise ValueError(f"Load DAG error, DAG ID {dag_id} has already exist")
+ triggers += dag.trigger_nodes
+ from ..trigger.trigger_manager import DefaultTriggerManager
+
+ trigger_manager: DefaultTriggerManager = self.system_app.get_component(
+ ComponentType.AWEL_TRIGGER_MANAGER,
+ DefaultTriggerManager,
+ default_component=None,
+ )
+ if trigger_manager:
+ for trigger in triggers:
+ trigger_manager.register_trigger(trigger)
+ trigger_manager.after_register()
+ else:
+ logger.warn("No trigger manager, not register dag trigger")
diff --git a/pilot/awel/dag/loader.py b/pilot/awel/dag/loader.py
new file mode 100644
index 000000000..2eb89f8bc
--- /dev/null
+++ b/pilot/awel/dag/loader.py
@@ -0,0 +1,93 @@
+from abc import ABC, abstractmethod
+from typing import List
+import os
+import hashlib
+import sys
+import logging
+import traceback
+
+from .base import DAG
+
+logger = logging.getLogger(__name__)
+
+
+class DAGLoader(ABC):
+ @abstractmethod
+ def load_dags(self) -> List[DAG]:
+ """Load dags"""
+
+
+class LocalFileDAGLoader(DAGLoader):
+ def __init__(self, filepath: str) -> None:
+ super().__init__()
+ self._filepath = filepath
+
+ def load_dags(self) -> List[DAG]:
+ if not os.path.exists(self._filepath):
+ return []
+ if os.path.isdir(self._filepath):
+ return _process_directory(self._filepath)
+ else:
+ return _process_file(self._filepath)
+
+
+def _process_directory(directory: str) -> List[DAG]:
+ dags = []
+ for file in os.listdir(directory):
+ if file.endswith(".py"):
+ filepath = os.path.join(directory, file)
+ dags += _process_file(filepath)
+ return dags
+
+
+def _process_file(filepath) -> List[DAG]:
+ mods = _load_modules_from_file(filepath)
+ results = _process_modules(mods)
+ return results
+
+
+def _load_modules_from_file(filepath: str):
+ import importlib
+ import importlib.machinery
+ import importlib.util
+
+ logger.info(f"Importing {filepath}")
+
+ org_mod_name, _ = os.path.splitext(os.path.split(filepath)[-1])
+ path_hash = hashlib.sha1(filepath.encode("utf-8")).hexdigest()
+ mod_name = f"unusual_prefix_{path_hash}_{org_mod_name}"
+
+ if mod_name in sys.modules:
+ del sys.modules[mod_name]
+
+ def parse(mod_name, filepath):
+ try:
+ loader = importlib.machinery.SourceFileLoader(mod_name, filepath)
+ spec = importlib.util.spec_from_loader(mod_name, loader)
+ new_module = importlib.util.module_from_spec(spec)
+ sys.modules[spec.name] = new_module
+ loader.exec_module(new_module)
+ return [new_module]
+ except Exception as e:
+ msg = traceback.format_exc()
+ logger.error(f"Failed to import: {filepath}, error message: {msg}")
+ # TODO save error message
+ return []
+
+ return parse(mod_name, filepath)
+
+
+def _process_modules(mods) -> List[DAG]:
+ top_level_dags = (
+ (o, m) for m in mods for o in m.__dict__.values() if isinstance(o, DAG)
+ )
+ found_dags = []
+ for dag, mod in top_level_dags:
+ try:
+ # TODO validate dag params
+ logger.info(f"Found dag {dag} from mod {mod} and model file {mod.__file__}")
+ found_dags.append(dag)
+ except Exception:
+ msg = traceback.format_exc()
+ logger.error(f"Failed to dag file, error message: {msg}")
+ return found_dags
diff --git a/pilot/awel/dag/tests/__init__.py b/pilot/awel/dag/tests/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/pilot/awel/dag/tests/test_dag.py b/pilot/awel/dag/tests/test_dag.py
new file mode 100644
index 000000000..c30530dc8
--- /dev/null
+++ b/pilot/awel/dag/tests/test_dag.py
@@ -0,0 +1,51 @@
+import pytest
+import threading
+import asyncio
+from ..dag import DAG, DAGContext
+
+
+def test_dag_context_sync():
+ dag1 = DAG("dag1")
+ dag2 = DAG("dag2")
+
+ with dag1:
+ assert DAGContext.get_current_dag() == dag1
+ with dag2:
+ assert DAGContext.get_current_dag() == dag2
+ assert DAGContext.get_current_dag() == dag1
+ assert DAGContext.get_current_dag() is None
+
+
+def test_dag_context_threading():
+ def thread_function(dag):
+ DAGContext.enter_dag(dag)
+ assert DAGContext.get_current_dag() == dag
+ DAGContext.exit_dag()
+
+ dag1 = DAG("dag1")
+ dag2 = DAG("dag2")
+
+ thread1 = threading.Thread(target=thread_function, args=(dag1,))
+ thread2 = threading.Thread(target=thread_function, args=(dag2,))
+
+ thread1.start()
+ thread2.start()
+ thread1.join()
+ thread2.join()
+
+ assert DAGContext.get_current_dag() is None
+
+
+@pytest.mark.asyncio
+async def test_dag_context_async():
+ async def async_function(dag):
+ DAGContext.enter_dag(dag)
+ assert DAGContext.get_current_dag() == dag
+ DAGContext.exit_dag()
+
+ dag1 = DAG("dag1")
+ dag2 = DAG("dag2")
+
+ await asyncio.gather(async_function(dag1), async_function(dag2))
+
+ assert DAGContext.get_current_dag() is None
diff --git a/pilot/awel/operator/__init__.py b/pilot/awel/operator/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/pilot/awel/operator/base.py b/pilot/awel/operator/base.py
new file mode 100644
index 000000000..09aa87141
--- /dev/null
+++ b/pilot/awel/operator/base.py
@@ -0,0 +1,206 @@
+from abc import ABC, abstractmethod, ABCMeta
+
+from types import FunctionType
+from typing import (
+ List,
+ Generic,
+ TypeVar,
+ AsyncIterator,
+ Union,
+ Any,
+ Dict,
+ Optional,
+ cast,
+)
+import functools
+from inspect import signature
+from pilot.component import SystemApp, ComponentType
+from pilot.utils.executor_utils import (
+ ExecutorFactory,
+ DefaultExecutorFactory,
+ blocking_func_to_async,
+ BlockingFunction,
+)
+
+from ..dag.base import DAGNode, DAGContext, DAGVar, DAG
+from ..task.base import (
+ TaskContext,
+ TaskOutput,
+ TaskState,
+ OUT,
+ T,
+ InputContext,
+ InputSource,
+)
+
+F = TypeVar("F", bound=FunctionType)
+
+CALL_DATA = Union[Dict, Dict[str, Dict]]
+
+
+class WorkflowRunner(ABC, Generic[T]):
+ """Abstract base class representing a runner for executing workflows in a DAG.
+
+ This class defines the interface for executing workflows within the DAG,
+ handling the flow from one DAG node to another.
+ """
+
+ @abstractmethod
+ async def execute_workflow(
+ self, node: "BaseOperator", call_data: Optional[CALL_DATA] = None
+ ) -> DAGContext:
+ """Execute the workflow starting from a given operator.
+
+ Args:
+ node (RunnableDAGNode): The starting node of the workflow to be executed.
+ call_data (CALL_DATA): The data pass to root operator node.
+
+ Returns:
+ DAGContext: The context after executing the workflow, containing the final state and data.
+ """
+
+
+default_runner: WorkflowRunner = None
+
+
+class BaseOperatorMeta(ABCMeta):
+ """Metaclass of BaseOperator."""
+
+ @classmethod
+ def _apply_defaults(cls, func: F) -> F:
+ sig_cache = signature(func)
+
+ @functools.wraps(func)
+ def apply_defaults(self: "BaseOperator", *args: Any, **kwargs: Any) -> Any:
+ dag: Optional[DAG] = kwargs.get("dag") or DAGVar.get_current_dag()
+ task_id: Optional[str] = kwargs.get("task_id")
+ system_app: Optional[SystemApp] = (
+ kwargs.get("system_app") or DAGVar.get_current_system_app()
+ )
+ executor = kwargs.get("executor") or DAGVar.get_executor()
+ if not executor:
+ if system_app:
+ executor = system_app.get_component(
+ ComponentType.EXECUTOR_DEFAULT, ExecutorFactory
+ ).create()
+ else:
+ executor = DefaultExecutorFactory().create()
+ DAGVar.set_executor(executor)
+
+ if not task_id and dag:
+ task_id = dag._new_node_id()
+ runner: Optional[WorkflowRunner] = kwargs.get("runner") or default_runner
+ # print(f"self: {self}, kwargs dag: {kwargs.get('dag')}, kwargs: {kwargs}")
+ # for arg in sig_cache.parameters:
+ # if arg not in kwargs:
+ # kwargs[arg] = default_args[arg]
+ if not kwargs.get("dag"):
+ kwargs["dag"] = dag
+ if not kwargs.get("task_id"):
+ kwargs["task_id"] = task_id
+ if not kwargs.get("runner"):
+ kwargs["runner"] = runner
+ if not kwargs.get("system_app"):
+ kwargs["system_app"] = system_app
+ if not kwargs.get("executor"):
+ kwargs["executor"] = executor
+ real_obj = func(self, *args, **kwargs)
+ return real_obj
+
+ return cast(T, apply_defaults)
+
+ def __new__(cls, name, bases, namespace, **kwargs):
+ new_cls = super().__new__(cls, name, bases, namespace, **kwargs)
+ new_cls.__init__ = cls._apply_defaults(new_cls.__init__)
+ return new_cls
+
+
+class BaseOperator(DAGNode, ABC, Generic[OUT], metaclass=BaseOperatorMeta):
+ """Abstract base class for operator nodes that can be executed within a workflow.
+
+ This class extends DAGNode by adding execution capabilities.
+ """
+
+ def __init__(
+ self,
+ task_id: Optional[str] = None,
+ task_name: Optional[str] = None,
+ dag: Optional[DAG] = None,
+ runner: WorkflowRunner = None,
+ **kwargs,
+ ) -> None:
+ """Initializes a BaseOperator with an optional workflow runner.
+
+ Args:
+ runner (WorkflowRunner, optional): The runner used to execute the workflow. Defaults to None.
+ """
+ super().__init__(node_id=task_id, node_name=task_name, dag=dag, **kwargs)
+ if not runner:
+ from pilot.awel import DefaultWorkflowRunner
+
+ runner = DefaultWorkflowRunner()
+
+ self._runner: WorkflowRunner = runner
+ self._dag_ctx: DAGContext = None
+
+ @property
+ def current_dag_context(self) -> DAGContext:
+ return self._dag_ctx
+
+ async def _run(self, dag_ctx: DAGContext) -> TaskOutput[OUT]:
+ if not self.node_id:
+ raise ValueError(f"The DAG Node ID can't be empty, current node {self}")
+ self._dag_ctx = dag_ctx
+ return await self._do_run(dag_ctx)
+
+ @abstractmethod
+ async def _do_run(self, dag_ctx: DAGContext) -> TaskOutput[OUT]:
+ """
+ Abstract method to run the task within the DAG node.
+
+ Args:
+ dag_ctx (DAGContext): The context of the DAG when this node is run.
+
+ Returns:
+ TaskOutput[OUT]: The task output after this node has been run.
+ """
+
+ async def call(self, call_data: Optional[CALL_DATA] = None) -> OUT:
+ """Execute the node and return the output.
+
+ This method is a high-level wrapper for executing the node.
+
+ Args:
+ call_data (CALL_DATA): The data pass to root operator node.
+
+ Returns:
+ OUT: The output of the node after execution.
+ """
+ out_ctx = await self._runner.execute_workflow(self, call_data)
+ return out_ctx.current_task_context.task_output.output
+
+ async def call_stream(
+ self, call_data: Optional[CALL_DATA] = None
+ ) -> AsyncIterator[OUT]:
+ """Execute the node and return the output as a stream.
+
+ This method is used for nodes where the output is a stream.
+
+ Args:
+ call_data (CALL_DATA): The data pass to root operator node.
+
+ Returns:
+ AsyncIterator[OUT]: An asynchronous iterator over the output stream.
+ """
+ out_ctx = await self._runner.execute_workflow(self, call_data)
+ return out_ctx.current_task_context.task_output.output_stream
+
+ async def blocking_func_to_async(
+ self, func: BlockingFunction, *args, **kwargs
+ ) -> Any:
+ return await blocking_func_to_async(self._executor, func, *args, **kwargs)
+
+
+def initialize_runner(runner: WorkflowRunner):
+ global default_runner
+ default_runner = runner
diff --git a/pilot/awel/operator/common_operator.py b/pilot/awel/operator/common_operator.py
new file mode 100644
index 000000000..2c0d41dde
--- /dev/null
+++ b/pilot/awel/operator/common_operator.py
@@ -0,0 +1,246 @@
+from typing import Generic, Dict, List, Union, Callable, Any, AsyncIterator, Awaitable
+import asyncio
+import logging
+
+from ..dag.base import DAGContext
+from ..task.base import (
+ TaskContext,
+ TaskOutput,
+ IN,
+ OUT,
+ InputContext,
+ InputSource,
+)
+
+from .base import BaseOperator
+
+
+logger = logging.getLogger(__name__)
+
+
+class JoinOperator(BaseOperator, Generic[OUT]):
+ """Operator that joins inputs using a custom combine function.
+
+ This node type is useful for combining the outputs of upstream nodes.
+ """
+
+ def __init__(self, combine_function, **kwargs):
+ super().__init__(**kwargs)
+ if not callable(combine_function):
+ raise ValueError("combine_function must be callable")
+ self.combine_function = combine_function
+
+ async def _do_run(self, dag_ctx: DAGContext) -> TaskOutput[OUT]:
+ """Run the join operation on the DAG context's inputs.
+ Args:
+ dag_ctx (DAGContext): The current context of the DAG.
+
+ Returns:
+ TaskOutput[OUT]: The task output after this node has been run.
+ """
+ curr_task_ctx: TaskContext[OUT] = dag_ctx.current_task_context
+ input_ctx: InputContext = await curr_task_ctx.task_input.map_all(
+ self.combine_function
+ )
+ # All join result store in the first parent output
+ join_output = input_ctx.parent_outputs[0].task_output
+ curr_task_ctx.set_task_output(join_output)
+ return join_output
+
+
+class ReduceStreamOperator(BaseOperator, Generic[IN, OUT]):
+ def __init__(self, reduce_function=None, **kwargs):
+ """Initializes a ReduceStreamOperator with a combine function.
+
+ Args:
+ combine_function: A function that defines how to combine inputs.
+
+ Raises:
+ ValueError: If the combine_function is not callable.
+ """
+ super().__init__(**kwargs)
+ if reduce_function and not callable(reduce_function):
+ raise ValueError("reduce_function must be callable")
+ self.reduce_function = reduce_function
+
+ async def _do_run(self, dag_ctx: DAGContext) -> TaskOutput[OUT]:
+ """Run the join operation on the DAG context's inputs.
+
+ Args:
+ dag_ctx (DAGContext): The current context of the DAG.
+
+ Returns:
+ TaskOutput[OUT]: The task output after this node has been run.
+ """
+ curr_task_ctx: TaskContext[OUT] = dag_ctx.current_task_context
+ task_input = curr_task_ctx.task_input
+ if not task_input.check_stream():
+ raise ValueError("ReduceStreamOperator expects stream data")
+ if not task_input.check_single_parent():
+ raise ValueError("ReduceStreamOperator expects single parent")
+
+ reduce_function = self.reduce_function or self.reduce
+
+ input_ctx: InputContext = await task_input.reduce(reduce_function)
+ # All join result store in the first parent output
+ reduce_output = input_ctx.parent_outputs[0].task_output
+ curr_task_ctx.set_task_output(reduce_output)
+ return reduce_output
+
+ async def reduce(self, input_value: AsyncIterator[IN]) -> OUT:
+ raise NotImplementedError
+
+
+class MapOperator(BaseOperator, Generic[IN, OUT]):
+ """Map operator that applies a mapping function to its inputs.
+
+ This operator transforms its input data using a provided mapping function and
+ passes the transformed data downstream.
+ """
+
+ def __init__(self, map_function=None, **kwargs):
+ """Initializes a MapDAGNode with a mapping function.
+
+ Args:
+ map_function: A function that defines how to map the input data.
+
+ Raises:
+ ValueError: If the map_function is not callable.
+ """
+ super().__init__(**kwargs)
+ if map_function and not callable(map_function):
+ raise ValueError("map_function must be callable")
+ self.map_function = map_function
+
+ async def _do_run(self, dag_ctx: DAGContext) -> TaskOutput[OUT]:
+ """Run the mapping operation on the DAG context's inputs.
+
+ This method applies the mapping function to the input context and updates
+ the DAG context with the new data.
+
+ Args:
+ dag_ctx (DAGContext[IN]): The current context of the DAG.
+
+ Returns:
+ TaskOutput[OUT]: The task output after this node has been run.
+
+ Raises:
+ ValueError: If not a single parent or the map_function is not callable
+ """
+ curr_task_ctx: TaskContext[OUT] = dag_ctx.current_task_context
+ if not curr_task_ctx.task_input.check_single_parent():
+ num_parents = len(curr_task_ctx.task_input.parent_outputs)
+ raise ValueError(
+ f"task {curr_task_ctx.task_id} MapDAGNode expects single parent, now number of parents: {num_parents}"
+ )
+ map_function = self.map_function or self.map
+
+ input_ctx: InputContext = await curr_task_ctx.task_input.map(map_function)
+ # All join result store in the first parent output
+ reduce_output = input_ctx.parent_outputs[0].task_output
+ curr_task_ctx.set_task_output(reduce_output)
+ return reduce_output
+
+ async def map(self, input_value: IN) -> OUT:
+ raise NotImplementedError
+
+
+BranchFunc = Union[Callable[[IN], bool], Callable[[IN], Awaitable[bool]]]
+
+
+class BranchOperator(BaseOperator, Generic[IN, OUT]):
+ """Operator node that branches the workflow based on a provided function.
+
+ This node filters its input data using a branching function and
+ allows for conditional paths in the workflow.
+ """
+
+ def __init__(
+ self, branches: Dict[BranchFunc[IN], Union[BaseOperator, str]], **kwargs
+ ):
+ """
+ Initializes a BranchDAGNode with a branching function.
+
+ Args:
+ branches (Dict[BranchFunc[IN], Union[BaseOperator, str]]): Dict of function that defines the branching condition.
+
+ Raises:
+ ValueError: If the branch_function is not callable.
+ """
+ super().__init__(**kwargs)
+ if branches:
+ for branch_function, value in branches.items():
+ if not callable(branch_function):
+ raise ValueError("branch_function must be callable")
+ if isinstance(value, BaseOperator):
+ branches[branch_function] = value.node_name or value.node_name
+ self._branches = branches
+
+ async def _do_run(self, dag_ctx: DAGContext) -> TaskOutput[OUT]:
+ """Run the branching operation on the DAG context's inputs.
+
+ This method applies the branching function to the input context to determine
+ the path of execution in the workflow.
+
+ Args:
+ dag_ctx (DAGContext[IN]): The current context of the DAG.
+
+ Returns:
+ TaskOutput[OUT]: The task output after this node has been run.
+ """
+ curr_task_ctx: TaskContext[OUT] = dag_ctx.current_task_context
+ task_input = curr_task_ctx.task_input
+ if task_input.check_stream():
+ raise ValueError("BranchDAGNode expects no stream data")
+ if not task_input.check_single_parent():
+ raise ValueError("BranchDAGNode expects single parent")
+
+ branches = self._branches
+ if not branches:
+ branches = await self.branchs()
+
+ branch_func_tasks = []
+ branch_nodes: List[str] = []
+ for func, node_name in branches.items():
+ branch_nodes.append(node_name)
+ branch_func_tasks.append(
+ curr_task_ctx.task_input.predicate_map(func, failed_value=None)
+ )
+
+ branch_input_ctxs: List[InputContext] = await asyncio.gather(*branch_func_tasks)
+ parent_output = task_input.parent_outputs[0].task_output
+ curr_task_ctx.set_task_output(parent_output)
+ skip_node_names = []
+ for i, ctx in enumerate(branch_input_ctxs):
+ node_name = branch_nodes[i]
+ branch_out = ctx.parent_outputs[0].task_output
+ logger.info(
+ f"branch_input_ctxs {i} result {branch_out.output}, is_empty: {branch_out.is_empty}"
+ )
+ if ctx.parent_outputs[0].task_output.is_empty:
+ logger.info(f"Skip node name {node_name}")
+ skip_node_names.append(node_name)
+ curr_task_ctx.update_metadata("skip_node_names", skip_node_names)
+ return parent_output
+
+ async def branchs(self) -> Dict[BranchFunc[IN], Union[BaseOperator, str]]:
+ raise NotImplementedError
+
+
+class InputOperator(BaseOperator, Generic[OUT]):
+ def __init__(self, input_source: InputSource[OUT], **kwargs) -> None:
+ super().__init__(**kwargs)
+ self._input_source = input_source
+
+ async def _do_run(self, dag_ctx: DAGContext) -> TaskOutput[OUT]:
+ curr_task_ctx: TaskContext[OUT] = dag_ctx.current_task_context
+ task_output = await self._input_source.read(curr_task_ctx)
+ curr_task_ctx.set_task_output(task_output)
+ return task_output
+
+
+class TriggerOperator(InputOperator, Generic[OUT]):
+ def __init__(self, **kwargs) -> None:
+ from ..task.task_impl import SimpleCallDataInputSource
+
+ super().__init__(input_source=SimpleCallDataInputSource(), **kwargs)
diff --git a/pilot/awel/operator/stream_operator.py b/pilot/awel/operator/stream_operator.py
new file mode 100644
index 000000000..7de916a83
--- /dev/null
+++ b/pilot/awel/operator/stream_operator.py
@@ -0,0 +1,90 @@
+from abc import ABC, abstractmethod
+from typing import Generic, AsyncIterator
+from ..task.base import OUT, IN, TaskOutput, TaskContext
+from ..dag.base import DAGContext
+from .base import BaseOperator
+
+
+class StreamifyAbsOperator(BaseOperator[OUT], ABC, Generic[IN, OUT]):
+ async def _do_run(self, dag_ctx: DAGContext) -> TaskOutput[OUT]:
+ curr_task_ctx: TaskContext[OUT] = dag_ctx.current_task_context
+ output = await curr_task_ctx.task_input.parent_outputs[0].task_output.streamify(
+ self.streamify
+ )
+ curr_task_ctx.set_task_output(output)
+ return output
+
+ @abstractmethod
+ async def streamify(self, input_value: IN) -> AsyncIterator[OUT]:
+ """Convert a value of IN to an AsyncIterator[OUT]
+
+ Args:
+ input_value (IN): The data of parent operator's output
+
+ Example:
+
+ .. code-block:: python
+
+ class MyStreamOperator(StreamifyAbsOperator[int, int]):
+ async def streamify(self, input_value: int) -> AsyncIterator[int]
+ for i in range(input_value):
+ yield i
+ """
+
+
+class UnstreamifyAbsOperator(BaseOperator[OUT], Generic[IN, OUT]):
+ async def _do_run(self, dag_ctx: DAGContext) -> TaskOutput[OUT]:
+ curr_task_ctx: TaskContext[OUT] = dag_ctx.current_task_context
+ output = await curr_task_ctx.task_input.parent_outputs[
+ 0
+ ].task_output.unstreamify(self.unstreamify)
+ curr_task_ctx.set_task_output(output)
+ return output
+
+ @abstractmethod
+ async def unstreamify(self, input_value: AsyncIterator[IN]) -> OUT:
+ """Convert a value of AsyncIterator[IN] to an OUT.
+
+ Args:
+ input_value (AsyncIterator[IN])): The data of parent operator's output
+
+ Example:
+
+ .. code-block:: python
+
+ class MyUnstreamOperator(UnstreamifyAbsOperator[int, int]):
+ async def unstreamify(self, input_value: AsyncIterator[int]) -> int
+ value_cnt = 0
+ async for v in input_value:
+ value_cnt += 1
+ return value_cnt
+ """
+
+
+class TransformStreamAbsOperator(BaseOperator[OUT], Generic[IN, OUT]):
+ async def _do_run(self, dag_ctx: DAGContext) -> TaskOutput[OUT]:
+ curr_task_ctx: TaskContext[OUT] = dag_ctx.current_task_context
+ output = await curr_task_ctx.task_input.parent_outputs[
+ 0
+ ].task_output.transform_stream(self.transform_stream)
+ curr_task_ctx.set_task_output(output)
+ return output
+
+ @abstractmethod
+ async def transform_stream(
+ self, input_value: AsyncIterator[IN]
+ ) -> AsyncIterator[OUT]:
+ """Transform an AsyncIterator[IN] to another AsyncIterator[OUT] using a given function.
+
+ Args:
+ input_value (AsyncIterator[IN])): The data of parent operator's output
+
+ Example:
+
+ .. code-block:: python
+
+ class MyTransformStreamOperator(TransformStreamAbsOperator[int, int]):
+ async def unstreamify(self, input_value: AsyncIterator[int]) -> AsyncIterator[int]
+ async for v in input_value:
+ yield v + 1
+ """
diff --git a/pilot/awel/resource/__init__.py b/pilot/awel/resource/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/pilot/awel/resource/base.py b/pilot/awel/resource/base.py
new file mode 100644
index 000000000..97fefbbc3
--- /dev/null
+++ b/pilot/awel/resource/base.py
@@ -0,0 +1,8 @@
+from abc import ABC, abstractmethod
+
+
+class ResourceGroup(ABC):
+ @property
+ @abstractmethod
+ def name(self) -> str:
+ """The name of current resource group"""
diff --git a/pilot/awel/runner/__init__.py b/pilot/awel/runner/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/pilot/awel/runner/job_manager.py b/pilot/awel/runner/job_manager.py
new file mode 100644
index 000000000..7a1d12ead
--- /dev/null
+++ b/pilot/awel/runner/job_manager.py
@@ -0,0 +1,82 @@
+from typing import List, Set, Optional, Dict
+import uuid
+import logging
+from ..dag.base import DAG
+
+from ..operator.base import BaseOperator, CALL_DATA
+
+logger = logging.getLogger(__name__)
+
+
+class DAGNodeInstance:
+ def __init__(self, node_instance: DAG) -> None:
+ pass
+
+
+class DAGInstance:
+ def __init__(self, dag: DAG) -> None:
+ self._dag = dag
+
+
+class JobManager:
+ def __init__(
+ self,
+ root_nodes: List[BaseOperator],
+ all_nodes: List[BaseOperator],
+ end_node: BaseOperator,
+ id2call_data: Dict[str, Dict],
+ ) -> None:
+ self._root_nodes = root_nodes
+ self._all_nodes = all_nodes
+ self._end_node = end_node
+ self._id2node_data = id2call_data
+
+ @staticmethod
+ def build_from_end_node(
+ end_node: BaseOperator, call_data: Optional[CALL_DATA] = None
+ ) -> "JobManager":
+ nodes = _build_from_end_node(end_node)
+ root_nodes = _get_root_nodes(nodes)
+ id2call_data = _save_call_data(root_nodes, call_data)
+ return JobManager(root_nodes, nodes, end_node, id2call_data)
+
+ def get_call_data_by_id(self, node_id: str) -> Optional[Dict]:
+ return self._id2node_data.get(node_id)
+
+
+def _save_call_data(
+ root_nodes: List[BaseOperator], call_data: CALL_DATA
+) -> Dict[str, Dict]:
+ id2call_data = {}
+ logger.debug(f"_save_call_data: {call_data}, root_nodes: {root_nodes}")
+ if not call_data:
+ return id2call_data
+ if len(root_nodes) == 1:
+ node = root_nodes[0]
+ logger.info(f"Save call data to node {node.node_id}, call_data: {call_data}")
+ id2call_data[node.node_id] = call_data
+ else:
+ for node in root_nodes:
+ node_id = node.node_id
+ logger.info(
+ f"Save call data to node {node.node_id}, call_data: {call_data.get(node_id)}"
+ )
+ id2call_data[node_id] = call_data.get(node_id)
+ return id2call_data
+
+
+def _build_from_end_node(end_node: BaseOperator) -> List[BaseOperator]:
+ nodes = []
+ if isinstance(end_node, BaseOperator):
+ task_id = end_node.node_id
+ if not task_id:
+ task_id = str(uuid.uuid4())
+ end_node.set_node_id(task_id)
+ nodes.append(end_node)
+ for node in end_node.upstream:
+ nodes += _build_from_end_node(node)
+ return nodes
+
+
+def _get_root_nodes(nodes: List[BaseOperator]) -> List[BaseOperator]:
+ return list(set(filter(lambda x: not x.upstream, nodes)))
diff --git a/pilot/awel/runner/local_runner.py b/pilot/awel/runner/local_runner.py
new file mode 100644
index 000000000..6f8a0a484
--- /dev/null
+++ b/pilot/awel/runner/local_runner.py
@@ -0,0 +1,106 @@
+from typing import Dict, Optional, Set, List
+import logging
+
+from ..dag.base import DAGContext
+from ..operator.base import WorkflowRunner, BaseOperator, CALL_DATA
+from ..operator.common_operator import BranchOperator, JoinOperator, TriggerOperator
+from ..task.base import TaskContext, TaskState
+from ..task.task_impl import DefaultInputContext, DefaultTaskContext, SimpleTaskOutput
+from .job_manager import JobManager
+
+logger = logging.getLogger(__name__)
+
+
+class DefaultWorkflowRunner(WorkflowRunner):
+ async def execute_workflow(
+ self, node: BaseOperator, call_data: Optional[CALL_DATA] = None
+ ) -> DAGContext:
+ # Create DAG context
+ dag_ctx = DAGContext()
+ job_manager = JobManager.build_from_end_node(node, call_data)
+ logger.info(
+ f"Begin run workflow from end operator, id: {node.node_id}, call_data: {call_data}"
+ )
+ dag = node.dag
+ # Save node output
+ node_outputs: Dict[str, TaskContext] = {}
+ skip_node_ids = set()
+ await self._execute_node(
+ job_manager, node, dag_ctx, node_outputs, skip_node_ids
+ )
+
+ return dag_ctx
+
+ async def _execute_node(
+ self,
+ job_manager: JobManager,
+ node: BaseOperator,
+ dag_ctx: DAGContext,
+ node_outputs: Dict[str, TaskContext],
+ skip_node_ids: Set[str],
+ ):
+ # Skip run node
+ if node.node_id in node_outputs:
+ return
+
+ # Run all upstream node
+ for upstream_node in node.upstream:
+ if isinstance(upstream_node, BaseOperator):
+ await self._execute_node(
+ job_manager, upstream_node, dag_ctx, node_outputs, skip_node_ids
+ )
+
+ inputs = [
+ node_outputs[upstream_node.node_id] for upstream_node in node.upstream
+ ]
+ input_ctx = DefaultInputContext(inputs)
+ task_ctx = DefaultTaskContext(node.node_id, TaskState.INIT, task_output=None)
+ task_ctx.set_call_data(job_manager.get_call_data_by_id(node.node_id))
+
+ task_ctx.set_task_input(input_ctx)
+ dag_ctx.set_current_task_context(task_ctx)
+ task_ctx.set_current_state(TaskState.RUNNING)
+
+ if node.node_id in skip_node_ids:
+ task_ctx.set_current_state(TaskState.SKIP)
+ task_ctx.set_task_output(SimpleTaskOutput(None))
+ node_outputs[node.node_id] = task_ctx
+ return
+ try:
+ logger.debug(
+ f"Begin run operator, node id: {node.node_id}, node name: {node.node_name}, cls: {node}"
+ )
+ await node._run(dag_ctx)
+ node_outputs[node.node_id] = dag_ctx.current_task_context
+ task_ctx.set_current_state(TaskState.SUCCESS)
+
+ if isinstance(node, BranchOperator):
+ skip_nodes = task_ctx.metadata.get("skip_node_names", [])
+ logger.debug(
+ f"Current is branch operator, skip node names: {skip_nodes}"
+ )
+ _skip_current_downstream_by_node_name(node, skip_nodes, skip_node_ids)
+ except Exception as e:
+ logger.info(f"Run operator {node.node_id} error, error message: {str(e)}")
+ task_ctx.set_current_state(TaskState.FAILED)
+ raise e
+
+
+def _skip_current_downstream_by_node_name(
+ branch_node: BranchOperator, skip_nodes: List[str], skip_node_ids: Set[str]
+):
+ if not skip_nodes:
+ return
+ for child in branch_node.downstream:
+ if child.node_name in skip_nodes:
+ logger.info(f"Skip node name {child.node_name}, node id {child.node_id}")
+ _skip_downstream_by_id(child, skip_node_ids)
+
+
+def _skip_downstream_by_id(node: BaseOperator, skip_node_ids: Set[str]):
+ if isinstance(node, JoinOperator):
+ # Not skip join node
+ return
+ skip_node_ids.add(node.node_id)
+ for child in node.downstream:
+ _skip_downstream_by_id(child, skip_node_ids)
diff --git a/pilot/awel/task/__init__.py b/pilot/awel/task/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/pilot/awel/task/base.py b/pilot/awel/task/base.py
new file mode 100644
index 000000000..88b0df343
--- /dev/null
+++ b/pilot/awel/task/base.py
@@ -0,0 +1,367 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import (
+ TypeVar,
+ Generic,
+ Optional,
+ AsyncIterator,
+ Union,
+ Callable,
+ Any,
+ Dict,
+ List,
+)
+
+IN = TypeVar("IN")
+OUT = TypeVar("OUT")
+T = TypeVar("T")
+
+
+class TaskState(str, Enum):
+ """Enumeration representing the state of a task in the workflow.
+
+ This Enum defines various states a task can be in during its lifecycle in the DAG.
+ """
+
+ INIT = "init" # Initial state of the task, not yet started
+ SKIP = "skip" # State indicating the task was skipped
+ RUNNING = "running" # State indicating the task is currently running
+ SUCCESS = "success" # State indicating the task completed successfully
+ FAILED = "failed" # State indicating the task failed during execution
+
+
+class TaskOutput(ABC, Generic[T]):
+ """Abstract base class representing the output of a task.
+
+ This class encapsulates the output of a task and provides methods to access the output data.
+ It can be subclassed to implement specific output behaviors.
+ """
+
+ @property
+ def is_stream(self) -> bool:
+ """Check if the output is a stream.
+
+ Returns:
+ bool: True if the output is a stream, False otherwise.
+ """
+ return False
+
+ @property
+ def is_empty(self) -> bool:
+ """Check if the output is empty.
+
+ Returns:
+ bool: True if the output is empty, False otherwise.
+ """
+ return False
+
+ @property
+ def output(self) -> Optional[T]:
+ """Return the output of the task.
+
+ Returns:
+ T: The output of the task. None if the output is empty.
+ """
+ raise NotImplementedError
+
+ @property
+ def output_stream(self) -> Optional[AsyncIterator[T]]:
+ """Return the output of the task as an asynchronous stream.
+
+ Returns:
+ AsyncIterator[T]: An asynchronous iterator over the output. None if the output is empty.
+ """
+ raise NotImplementedError
+
+ @abstractmethod
+ def set_output(self, output_data: Union[T, AsyncIterator[T]]) -> None:
+ """Set the output data to current object.
+
+ Args:
+ output_data (Union[T, AsyncIterator[T]]): Output data.
+ """
+
+ @abstractmethod
+ def new_output(self) -> "TaskOutput[T]":
+ """Create new output object"""
+
+ async def map(self, map_func) -> "TaskOutput[T]":
+ """Apply a mapping function to the task's output.
+
+ Args:
+ map_func: A function to apply to the task's output.
+
+ Returns:
+ TaskOutput[T]: The result of applying the mapping function.
+ """
+ raise NotImplementedError
+
+ async def reduce(self, reduce_func) -> "TaskOutput[T]":
+ """Apply a reducing function to the task's output.
+
+ Stream TaskOutput to Nonstream TaskOutput.
+
+ Args:
+ reduce_func: A reducing function to apply to the task's output.
+
+ Returns:
+ TaskOutput[T]: The result of applying the reducing function.
+ """
+ raise NotImplementedError
+
+ async def streamify(
+ self, transform_func: Callable[[T], AsyncIterator[T]]
+ ) -> "TaskOutput[T]":
+ """Convert a value of type T to an AsyncIterator[T] using a transform function.
+
+ Args:
+ transform_func (Callable[[T], AsyncIterator[T]]): Function to transform a T value into an AsyncIterator[T].
+
+ Returns:
+ TaskOutput[T]: The result of applying the reducing function.
+ """
+ raise NotImplementedError
+
+ async def transform_stream(
+ self, transform_func: Callable[[AsyncIterator[T]], AsyncIterator[T]]
+ ) -> "TaskOutput[T]":
+ """Transform an AsyncIterator[T] to another AsyncIterator[T] using a given function.
+
+ Args:
+ transform_func (Callable[[AsyncIterator[T]], AsyncIterator[T]]): Function to apply to the AsyncIterator[T].
+
+ Returns:
+ TaskOutput[T]: The result of applying the reducing function.
+ """
+ raise NotImplementedError
+
+ async def unstreamify(
+ self, transform_func: Callable[[AsyncIterator[T]], T]
+ ) -> "TaskOutput[T]":
+ """Convert an AsyncIterator[T] to a value of type T using a transform function.
+
+ Args:
+ transform_func (Callable[[AsyncIterator[T]], T]): Function to transform an AsyncIterator[T] into a T value.
+
+ Returns:
+ TaskOutput[T]: The result of applying the reducing function.
+ """
+ raise NotImplementedError
+
+ async def check_condition(self, condition_func) -> bool:
+ """Check if current output meets a given condition.
+
+ Args:
+ condition_func: A function to determine if the condition is met.
+ Returns:
+ bool: True if current output meet the condition, False otherwise.
+ """
+ raise NotImplementedError
+
+
+class TaskContext(ABC, Generic[T]):
+ """Abstract base class representing the context of a task within a DAG.
+
+ This class provides the interface for accessing task-related information
+ and manipulating task output.
+ """
+
+ @property
+ @abstractmethod
+ def task_id(self) -> str:
+ """Return the unique identifier of the task.
+
+ Returns:
+ str: The unique identifier of the task.
+ """
+
+ @property
+ @abstractmethod
+ def task_input(self) -> "InputContext":
+ """Return the InputContext of current task.
+
+ Returns:
+ InputContext: The InputContext of current task.
+ """
+
+ @abstractmethod
+ def set_task_input(self, input_ctx: "InputContext") -> None:
+ """Set the InputContext object to current task.
+
+ Args:
+ input_ctx (InputContext): The InputContext of current task
+ """
+
+ @property
+ @abstractmethod
+ def task_output(self) -> TaskOutput[T]:
+ """Return the output object of the task.
+
+ Returns:
+ TaskOutput[T]: The output object of the task.
+ """
+
+ @abstractmethod
+ def set_task_output(self, task_output: TaskOutput[T]) -> None:
+ """Set the output object to current task."""
+
+ @property
+ @abstractmethod
+ def current_state(self) -> TaskState:
+ """Get the current state of the task.
+
+ Returns:
+ TaskState: The current state of the task.
+ """
+
+ @abstractmethod
+ def set_current_state(self, task_state: TaskState) -> None:
+ """Set current task state
+
+ Args:
+ task_state (TaskState): The task state to be set.
+ """
+
+ @abstractmethod
+ def new_ctx(self) -> "TaskContext":
+ """Create new task context
+
+ Returns:
+ TaskContext: A new instance of a TaskContext.
+ """
+
+ @property
+ @abstractmethod
+ def metadata(self) -> Dict[str, Any]:
+ """Get the metadata of current task
+
+ Returns:
+ Dict[str, Any]: The metadata
+ """
+
+ def update_metadata(self, key: str, value: Any) -> None:
+ """Update metadata with key and value
+
+ Args:
+ key (str): The key of metadata
+ value (str): The value to be add to metadata
+ """
+ self.metadata[key] = value
+
+ @property
+ def call_data(self) -> Optional[Dict]:
+ """Get the call data for current data"""
+ return self.metadata.get("call_data")
+
+ def set_call_data(self, call_data: Dict) -> None:
+ """Set call data for current task"""
+ self.update_metadata("call_data", call_data)
+
+
+class InputContext(ABC):
+ """Abstract base class representing the context of inputs to a operator node.
+
+ This class defines methods to manipulate and access the inputs for a operator node.
+ """
+
+ @property
+ @abstractmethod
+ def parent_outputs(self) -> List[TaskContext]:
+ """Get the outputs from the parent nodes.
+
+ Returns:
+ List[TaskContext]: A list of contexts of the parent nodes' outputs.
+ """
+
+ @abstractmethod
+ async def map(self, map_func: Callable[[Any], Any]) -> "InputContext":
+ """Apply a mapping function to the inputs.
+
+ Args:
+ map_func (Callable[[Any], Any]): A function to be applied to the inputs.
+
+ Returns:
+ InputContext: A new InputContext instance with the mapped inputs.
+ """
+
+ @abstractmethod
+ async def map_all(self, map_func: Callable[..., Any]) -> "InputContext":
+ """Apply a mapping function to all inputs.
+
+ Args:
+ map_func (Callable[..., Any]): A function to be applied to all inputs.
+
+ Returns:
+ InputContext: A new InputContext instance with the mapped inputs.
+ """
+
+ @abstractmethod
+ async def reduce(self, reduce_func: Callable[[Any], Any]) -> "InputContext":
+ """Apply a reducing function to the inputs.
+
+ Args:
+ reduce_func (Callable[[Any], Any]): A function that reduces the inputs.
+
+ Returns:
+ InputContext: A new InputContext instance with the reduced inputs.
+ """
+
+ @abstractmethod
+ async def filter(self, filter_func: Callable[[Any], bool]) -> "InputContext":
+ """Filter the inputs based on a provided function.
+
+ Args:
+ filter_func (Callable[[Any], bool]): A function that returns True for inputs to keep.
+
+ Returns:
+ InputContext: A new InputContext instance with the filtered inputs.
+ """
+
+ @abstractmethod
+ async def predicate_map(
+ self, predicate_func: Callable[[Any], bool], failed_value: Any = None
+ ) -> "InputContext":
+ """Predicate the inputs based on a provided function.
+
+ Args:
+ predicate_func (Callable[[Any], bool]): A function that returns True for inputs is predicate True.
+ failed_value (Any): The value to be set if the return value of predicate function is False
+ Returns:
+ InputContext: A new InputContext instance with the predicate inputs.
+ """
+
+ def check_single_parent(self) -> bool:
+ """Check if there is only a single parent output.
+
+ Returns:
+ bool: True if there is only one parent output, False otherwise.
+ """
+ return len(self.parent_outputs) == 1
+
+ def check_stream(self, skip_empty: bool = False) -> bool:
+ """Check if all parent outputs are streams.
+
+ Args:
+ skip_empty (bool): Skip empty output or not.
+
+ Returns:
+ bool: True if all parent outputs are streams, False otherwise.
+ """
+ for out in self.parent_outputs:
+ if out.task_output.is_empty and skip_empty:
+ continue
+ if not (out.task_output and out.task_output.is_stream):
+ return False
+ return True
+
+
+class InputSource(ABC, Generic[T]):
+ """Abstract base class representing the source of inputs to a DAG node."""
+
+ @abstractmethod
+ async def read(self, task_ctx: TaskContext) -> TaskOutput[T]:
+ """Read the data from current input source.
+
+ Returns:
+ TaskOutput[T]: The output object read from current source
+ """
diff --git a/pilot/awel/task/task_impl.py b/pilot/awel/task/task_impl.py
new file mode 100644
index 000000000..f969c135c
--- /dev/null
+++ b/pilot/awel/task/task_impl.py
@@ -0,0 +1,339 @@
+from abc import ABC, abstractmethod
+from typing import (
+ Callable,
+ Coroutine,
+ Iterator,
+ AsyncIterator,
+ List,
+ Generic,
+ TypeVar,
+ Any,
+ Tuple,
+ Dict,
+ Union,
+)
+import asyncio
+import logging
+from .base import TaskOutput, TaskContext, TaskState, InputContext, InputSource, T
+
+
+logger = logging.getLogger(__name__)
+
+
+async def _reduce_stream(stream: AsyncIterator, reduce_function) -> Any:
+ # Init accumulator
+ try:
+ accumulator = await stream.__anext__()
+ except StopAsyncIteration:
+ raise ValueError("Stream is empty")
+ is_async = asyncio.iscoroutinefunction(reduce_function)
+ async for element in stream:
+ if is_async:
+ accumulator = await reduce_function(accumulator, element)
+ else:
+ accumulator = reduce_function(accumulator, element)
+ return accumulator
+
+
+class SimpleTaskOutput(TaskOutput[T], Generic[T]):
+ def __init__(self, data: T) -> None:
+ super().__init__()
+ self._data = data
+
+ @property
+ def output(self) -> T:
+ return self._data
+
+ def set_output(self, output_data: T | AsyncIterator[T]) -> None:
+ self._data = output_data
+
+ def new_output(self) -> TaskOutput[T]:
+ return SimpleTaskOutput(None)
+
+ @property
+ def is_empty(self) -> bool:
+ return not self._data
+
+ async def _apply_func(self, func) -> Any:
+ if asyncio.iscoroutinefunction(func):
+ out = await func(self._data)
+ else:
+ out = func(self._data)
+ return out
+
+ async def map(self, map_func) -> TaskOutput[T]:
+ out = await self._apply_func(map_func)
+ return SimpleTaskOutput(out)
+
+ async def check_condition(self, condition_func) -> bool:
+ return await self._apply_func(condition_func)
+
+ async def streamify(
+ self, transform_func: Callable[[T], AsyncIterator[T]]
+ ) -> TaskOutput[T]:
+ out = await self._apply_func(transform_func)
+ return SimpleStreamTaskOutput(out)
+
+
+class SimpleStreamTaskOutput(TaskOutput[T], Generic[T]):
+ def __init__(self, data: AsyncIterator[T]) -> None:
+ super().__init__()
+ self._data = data
+
+ @property
+ def is_stream(self) -> bool:
+ return True
+
+ @property
+ def is_empty(self) -> bool:
+ return not self._data
+
+ @property
+ def output_stream(self) -> AsyncIterator[T]:
+ return self._data
+
+ def set_output(self, output_data: T | AsyncIterator[T]) -> None:
+ self._data = output_data
+
+ def new_output(self) -> TaskOutput[T]:
+ return SimpleStreamTaskOutput(None)
+
+ async def map(self, map_func) -> TaskOutput[T]:
+ is_async = asyncio.iscoroutinefunction(map_func)
+
+ async def new_iter() -> AsyncIterator[T]:
+ async for out in self._data:
+ if is_async:
+ out = await map_func(out)
+ else:
+ out = map_func(out)
+ yield out
+
+ return SimpleStreamTaskOutput(new_iter())
+
+ async def reduce(self, reduce_func) -> TaskOutput[T]:
+ out = await _reduce_stream(self._data, reduce_func)
+ return SimpleTaskOutput(out)
+
+ async def unstreamify(
+ self, transform_func: Callable[[AsyncIterator[T]], T]
+ ) -> TaskOutput[T]:
+ if asyncio.iscoroutinefunction(transform_func):
+ out = await transform_func(self._data)
+ else:
+ out = transform_func(self._data)
+ return SimpleTaskOutput(out)
+
+ async def transform_stream(
+ self, transform_func: Callable[[AsyncIterator[T]], AsyncIterator[T]]
+ ) -> TaskOutput[T]:
+ if asyncio.iscoroutinefunction(transform_func):
+ out = await transform_func(self._data)
+ else:
+ out = transform_func(self._data)
+ return SimpleStreamTaskOutput(out)
+
+
+def _is_async_iterator(obj):
+ return (
+ hasattr(obj, "__anext__")
+ and callable(getattr(obj, "__anext__", None))
+ and hasattr(obj, "__aiter__")
+ and callable(getattr(obj, "__aiter__", None))
+ )
+
+
+class BaseInputSource(InputSource, ABC):
+ def __init__(self) -> None:
+ super().__init__()
+ self._is_read = False
+
+ @abstractmethod
+ def _read_data(self, task_ctx: TaskContext) -> Any:
+ """Read data with task context"""
+
+ async def read(self, task_ctx: TaskContext) -> Coroutine[Any, Any, TaskOutput]:
+ data = self._read_data(task_ctx)
+ if _is_async_iterator(data):
+ if self._is_read:
+ raise ValueError(f"Input iterator {data} has been read!")
+ output = SimpleStreamTaskOutput(data)
+ else:
+ output = SimpleTaskOutput(data)
+ self._is_read = True
+ return output
+
+
+class SimpleInputSource(BaseInputSource):
+ def __init__(self, data: Any) -> None:
+ super().__init__()
+ self._data = data
+
+ def _read_data(self, task_ctx: TaskContext) -> Any:
+ return self._data
+
+
+class SimpleCallDataInputSource(BaseInputSource):
+ def __init__(self) -> None:
+ super().__init__()
+
+ def _read_data(self, task_ctx: TaskContext) -> Any:
+ call_data = task_ctx.call_data
+ data = call_data.get("data") if call_data else None
+ if not (call_data and data):
+ raise ValueError("No call data for current SimpleCallDataInputSource")
+ return data
+
+
+class DefaultTaskContext(TaskContext, Generic[T]):
+ def __init__(
+ self, task_id: str, task_state: TaskState, task_output: TaskOutput[T]
+ ) -> None:
+ super().__init__()
+ self._task_id = task_id
+ self._task_state = task_state
+ self._output = task_output
+ self._task_input = None
+ self._metadata = {}
+
+ @property
+ def task_id(self) -> str:
+ return self._task_id
+
+ @property
+ def task_input(self) -> InputContext:
+ return self._task_input
+
+ def set_task_input(self, input_ctx: "InputContext") -> None:
+ self._task_input = input_ctx
+
+ @property
+ def task_output(self) -> TaskOutput:
+ return self._output
+
+ def set_task_output(self, task_output: TaskOutput) -> None:
+ self._output = task_output
+
+ @property
+ def current_state(self) -> TaskState:
+ return self._task_state
+
+ def set_current_state(self, task_state: TaskState) -> None:
+ self._task_state = task_state
+
+ def new_ctx(self) -> TaskContext:
+ new_output = self._output.new_output()
+ return DefaultTaskContext(self._task_id, self._task_state, new_output)
+
+ @property
+ def metadata(self) -> Dict[str, Any]:
+ return self._metadata
+
+
+class DefaultInputContext(InputContext):
+ def __init__(self, outputs: List[TaskContext]) -> None:
+ super().__init__()
+ self._outputs = outputs
+
+ @property
+ def parent_outputs(self) -> List[TaskContext]:
+ return self._outputs
+
+ async def _apply_func(
+ self, func: Callable[[Any], Any], apply_type: str = "map"
+ ) -> Tuple[List[TaskContext], List[TaskOutput]]:
+ new_outputs: List[TaskContext] = []
+ map_tasks = []
+ for out in self._outputs:
+ new_outputs.append(out.new_ctx())
+ result = None
+ if apply_type == "map":
+ result = out.task_output.map(func)
+ elif apply_type == "reduce":
+ result = out.task_output.reduce(func)
+ elif apply_type == "check_condition":
+ result = out.task_output.check_condition(func)
+ else:
+ raise ValueError(f"Unsupport apply type {apply_type}")
+ map_tasks.append(result)
+ results = await asyncio.gather(*map_tasks)
+ return new_outputs, results
+
+ async def map(self, map_func: Callable[[Any], Any]) -> InputContext:
+ new_outputs, results = await self._apply_func(map_func)
+ for i, task_ctx in enumerate(new_outputs):
+ task_ctx: TaskContext = task_ctx
+ task_ctx.set_task_output(results[i])
+ return DefaultInputContext(new_outputs)
+
+ async def map_all(self, map_func: Callable[..., Any]) -> InputContext:
+ if not self._outputs:
+ return DefaultInputContext([])
+ # Some parent may be empty
+ not_empty_idx = 0
+ for i, p in enumerate(self._outputs):
+ if p.task_output.is_empty:
+ continue
+ not_empty_idx = i
+ break
+ # All output is empty?
+ is_steam = self._outputs[not_empty_idx].task_output.is_stream
+ if is_steam:
+ if not self.check_stream(skip_empty=True):
+ raise ValueError(
+ "The output in all tasks must has same output format to map_all"
+ )
+ outputs = []
+ for out in self._outputs:
+ if out.task_output.is_stream:
+ outputs.append(out.task_output.output_stream)
+ else:
+ outputs.append(out.task_output.output)
+ if asyncio.iscoroutinefunction(map_func):
+ map_res = await map_func(*outputs)
+ else:
+ map_res = map_func(*outputs)
+ single_output: TaskContext = self._outputs[not_empty_idx].new_ctx()
+ single_output.task_output.set_output(map_res)
+ logger.debug(
+ f"Current map_all map_res: {map_res}, is steam: {single_output.task_output.is_stream}"
+ )
+ return DefaultInputContext([single_output])
+
+ async def reduce(self, reduce_func: Callable[[Any], Any]) -> InputContext:
+ if not self.check_stream():
+ raise ValueError(
+ "The output in all tasks must has same output format of stream to apply reduce function"
+ )
+ new_outputs, results = await self._apply_func(reduce_func, apply_type="reduce")
+ for i, task_ctx in enumerate(new_outputs):
+ task_ctx: TaskContext = task_ctx
+ task_ctx.set_task_output(results[i])
+ return DefaultInputContext(new_outputs)
+
+ async def filter(self, filter_func: Callable[[Any], bool]) -> InputContext:
+ new_outputs, results = await self._apply_func(
+ filter_func, apply_type="check_condition"
+ )
+ result_outputs = []
+ for i, task_ctx in enumerate(new_outputs):
+ if results[i]:
+ result_outputs.append(task_ctx)
+ return DefaultInputContext(result_outputs)
+
+ async def predicate_map(
+ self, predicate_func: Callable[[Any], bool], failed_value: Any = None
+ ) -> "InputContext":
+ new_outputs, results = await self._apply_func(
+ predicate_func, apply_type="check_condition"
+ )
+ result_outputs = []
+ for i, task_ctx in enumerate(new_outputs):
+ task_ctx: TaskContext = task_ctx
+ if results[i]:
+ task_ctx.task_output.set_output(True)
+ result_outputs.append(task_ctx)
+ else:
+ task_ctx.task_output.set_output(failed_value)
+ result_outputs.append(task_ctx)
+ return DefaultInputContext(result_outputs)
diff --git a/pilot/awel/tests/__init__.py b/pilot/awel/tests/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/pilot/awel/tests/conftest.py b/pilot/awel/tests/conftest.py
new file mode 100644
index 000000000..2279cceba
--- /dev/null
+++ b/pilot/awel/tests/conftest.py
@@ -0,0 +1,102 @@
+import pytest
+import pytest_asyncio
+from typing import AsyncIterator, List
+from contextlib import contextmanager, asynccontextmanager
+from .. import (
+ WorkflowRunner,
+ InputOperator,
+ DAGContext,
+ TaskState,
+ DefaultWorkflowRunner,
+ SimpleInputSource,
+)
+from ..task.task_impl import _is_async_iterator
+
+
+@pytest.fixture
+def runner():
+ return DefaultWorkflowRunner()
+
+
+def _create_stream(num_nodes) -> List[AsyncIterator[int]]:
+ iters = []
+ for _ in range(num_nodes):
+
+ async def stream_iter():
+ for i in range(10):
+ yield i
+
+ stream_iter = stream_iter()
+ assert _is_async_iterator(stream_iter)
+ iters.append(stream_iter)
+ return iters
+
+
+def _create_stream_from(output_streams: List[List[int]]) -> List[AsyncIterator[int]]:
+ iters = []
+ for single_stream in output_streams:
+
+ async def stream_iter():
+ for i in single_stream:
+ yield i
+
+ stream_iter = stream_iter()
+ assert _is_async_iterator(stream_iter)
+ iters.append(stream_iter)
+ return iters
+
+
+@asynccontextmanager
+async def _create_input_node(**kwargs):
+ num_nodes = kwargs.get("num_nodes")
+ is_stream = kwargs.get("is_stream", False)
+ if is_stream:
+ outputs = kwargs.get("output_streams")
+ if outputs:
+ if num_nodes and num_nodes != len(outputs):
+ raise ValueError(
+ f"num_nodes {num_nodes} != the length of output_streams {len(outputs)}"
+ )
+ outputs = _create_stream_from(outputs)
+ else:
+ num_nodes = num_nodes or 1
+ outputs = _create_stream(num_nodes)
+ else:
+ outputs = kwargs.get("outputs", ["Hello."])
+ nodes = []
+ for output in outputs:
+ print(f"output: {output}")
+ input_source = SimpleInputSource(output)
+ input_node = InputOperator(input_source)
+ nodes.append(input_node)
+ yield nodes
+
+
+@pytest_asyncio.fixture
+async def input_node(request):
+ param = getattr(request, "param", {})
+ async with _create_input_node(**param) as input_nodes:
+ yield input_nodes[0]
+
+
+@pytest_asyncio.fixture
+async def stream_input_node(request):
+ param = getattr(request, "param", {})
+ param["is_stream"] = True
+ async with _create_input_node(**param) as input_nodes:
+ yield input_nodes[0]
+
+
+@pytest_asyncio.fixture
+async def input_nodes(request):
+ param = getattr(request, "param", {})
+ async with _create_input_node(**param) as input_nodes:
+ yield input_nodes
+
+
+@pytest_asyncio.fixture
+async def stream_input_nodes(request):
+ param = getattr(request, "param", {})
+ param["is_stream"] = True
+ async with _create_input_node(**param) as input_nodes:
+ yield input_nodes
diff --git a/pilot/awel/tests/test_http_operator.py b/pilot/awel/tests/test_http_operator.py
new file mode 100644
index 000000000..c57e70fe1
--- /dev/null
+++ b/pilot/awel/tests/test_http_operator.py
@@ -0,0 +1,51 @@
+import pytest
+from typing import List
+from .. import (
+ DAG,
+ WorkflowRunner,
+ DAGContext,
+ TaskState,
+ InputOperator,
+ MapOperator,
+ JoinOperator,
+ BranchOperator,
+ ReduceStreamOperator,
+ SimpleInputSource,
+)
+from .conftest import (
+ runner,
+ input_node,
+ input_nodes,
+ stream_input_node,
+ stream_input_nodes,
+ _is_async_iterator,
+)
+
+
+def _register_dag_to_fastapi_app(dag):
+ # TODO
+ pass
+
+
+@pytest.mark.asyncio
+async def test_http_operator(runner: WorkflowRunner, stream_input_node: InputOperator):
+ with DAG("test_map") as dag:
+ pass
+ # http_req_task = HttpRequestOperator(endpoint="/api/completions")
+ # db_task = DBQueryOperator(table_name="user_info")
+ # prompt_task = PromptTemplateOperator(
+ # system_prompt="You are an AI designed to solve the user's goals with given commands, please follow the constraints of the system's input for your answers."
+ # )
+ # llm_task = ChatGPTLLMOperator(model="chagpt-3.5")
+ # output_parser_task = CommonOutputParserOperator()
+ # http_res_task = HttpResponseOperator()
+ # (
+ # http_req_task
+ # >> db_task
+ # >> prompt_task
+ # >> llm_task
+ # >> output_parser_task
+ # >> http_res_task
+ # )
+
+ _register_dag_to_fastapi_app(dag)
diff --git a/pilot/awel/tests/test_run_dag.py b/pilot/awel/tests/test_run_dag.py
new file mode 100644
index 000000000..c0ea8e7ad
--- /dev/null
+++ b/pilot/awel/tests/test_run_dag.py
@@ -0,0 +1,141 @@
+import pytest
+from typing import List
+from .. import (
+ DAG,
+ WorkflowRunner,
+ DAGContext,
+ TaskState,
+ InputOperator,
+ MapOperator,
+ JoinOperator,
+ BranchOperator,
+ ReduceStreamOperator,
+ SimpleInputSource,
+)
+from .conftest import (
+ runner,
+ input_node,
+ input_nodes,
+ stream_input_node,
+ stream_input_nodes,
+ _is_async_iterator,
+)
+
+
+@pytest.mark.asyncio
+async def test_input_node(runner: WorkflowRunner):
+ input_node = InputOperator(SimpleInputSource("hello"))
+ res: DAGContext[str] = await runner.execute_workflow(input_node)
+ assert res.current_task_context.current_state == TaskState.SUCCESS
+ assert res.current_task_context.task_output.output == "hello"
+
+ async def new_steam_iter(n: int):
+ for i in range(n):
+ yield i
+
+ num_iter = 10
+ steam_input_node = InputOperator(SimpleInputSource(new_steam_iter(num_iter)))
+ res: DAGContext[str] = await runner.execute_workflow(steam_input_node)
+ assert res.current_task_context.current_state == TaskState.SUCCESS
+ output_steam = res.current_task_context.task_output.output_stream
+ assert output_steam
+ assert _is_async_iterator(output_steam)
+ i = 0
+ async for x in output_steam:
+ assert x == i
+ i += 1
+
+
+@pytest.mark.asyncio
+async def test_map_node(runner: WorkflowRunner, stream_input_node: InputOperator):
+ with DAG("test_map") as dag:
+ map_node = MapOperator(lambda x: x * 2)
+ stream_input_node >> map_node
+ res: DAGContext[int] = await runner.execute_workflow(map_node)
+ output_steam = res.current_task_context.task_output.output_stream
+ assert output_steam
+ i = 0
+ async for x in output_steam:
+ assert x == i * 2
+ i += 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "stream_input_node, expect_sum",
+ [
+ ({"output_streams": [[0, 1, 2, 3]]}, 6),
+ ({"output_streams": [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]}, 55),
+ ],
+ indirect=["stream_input_node"],
+)
+async def test_reduce_node(
+ runner: WorkflowRunner, stream_input_node: InputOperator, expect_sum: int
+):
+ with DAG("test_reduce_node") as dag:
+ reduce_node = ReduceStreamOperator(lambda x, y: x + y)
+ stream_input_node >> reduce_node
+ res: DAGContext[int] = await runner.execute_workflow(reduce_node)
+ assert res.current_task_context.current_state == TaskState.SUCCESS
+ assert not res.current_task_context.task_output.is_stream
+ assert res.current_task_context.task_output.output == expect_sum
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "input_nodes",
+ [
+ ({"outputs": [0, 1, 2]}),
+ ],
+ indirect=["input_nodes"],
+)
+async def test_join_node(runner: WorkflowRunner, input_nodes: List[InputOperator]):
+ def join_func(p1, p2, p3) -> int:
+ return p1 + p2 + p3
+
+ with DAG("test_join_node") as dag:
+ join_node = JoinOperator(join_func)
+ for input_node in input_nodes:
+ input_node >> join_node
+ res: DAGContext[int] = await runner.execute_workflow(join_node)
+ assert res.current_task_context.current_state == TaskState.SUCCESS
+ assert not res.current_task_context.task_output.is_stream
+ assert res.current_task_context.task_output.output == 3
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "input_node, is_odd",
+ [
+ ({"outputs": [0]}, False),
+ ({"outputs": [1]}, True),
+ ],
+ indirect=["input_node"],
+)
+async def test_branch_node(
+ runner: WorkflowRunner, input_node: InputOperator, is_odd: bool
+):
+ def join_func(o1, o2) -> int:
+ print(f"join func result, o1: {o1}, o2: {o2}")
+ return o1 or o2
+
+ with DAG("test_join_node") as dag:
+ odd_node = MapOperator(
+ lambda x: 999, task_id="odd_node", task_name="odd_node_name"
+ )
+ even_node = MapOperator(
+ lambda x: 888, task_id="even_node", task_name="even_node_name"
+ )
+ join_node = JoinOperator(join_func)
+ branch_node = BranchOperator(
+ {lambda x: x % 2 == 1: odd_node, lambda x: x % 2 == 0: even_node}
+ )
+ branch_node >> odd_node >> join_node
+ branch_node >> even_node >> join_node
+
+ input_node >> branch_node
+
+ res: DAGContext[int] = await runner.execute_workflow(join_node)
+ assert res.current_task_context.current_state == TaskState.SUCCESS
+ expect_res = 999 if is_odd else 888
+ assert res.current_task_context.task_output.output == expect_res
diff --git a/pilot/awel/trigger/__init__.py b/pilot/awel/trigger/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/pilot/awel/trigger/base.py b/pilot/awel/trigger/base.py
new file mode 100644
index 000000000..28662498f
--- /dev/null
+++ b/pilot/awel/trigger/base.py
@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from ..operator.common_operator import TriggerOperator
+
+
+class Trigger(TriggerOperator, ABC):
+ @abstractmethod
+ async def trigger(self) -> None:
+ """Trigger the workflow or a specific operation in the workflow."""
diff --git a/pilot/awel/trigger/http_trigger.py b/pilot/awel/trigger/http_trigger.py
new file mode 100644
index 000000000..175d4b63f
--- /dev/null
+++ b/pilot/awel/trigger/http_trigger.py
@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+from typing import Union, Type, List, TYPE_CHECKING, Optional, Any, Dict
+from starlette.requests import Request
+from starlette.responses import Response
+from pydantic import BaseModel
+import logging
+
+from .base import Trigger
+from ..dag.base import DAG
+from ..operator.base import BaseOperator
+
+if TYPE_CHECKING:
+ from fastapi import APIRouter, FastAPI
+
+RequestBody = Union[Request, Type[BaseModel], str]
+
+logger = logging.getLogger(__name__)
+
+
+class HttpTrigger(Trigger):
+ def __init__(
+ self,
+ endpoint: str,
+ methods: Optional[Union[str, List[str]]] = "GET",
+ request_body: Optional[RequestBody] = None,
+ streaming_response: Optional[bool] = False,
+ response_model: Optional[Type] = None,
+ response_headers: Optional[Dict[str, str]] = None,
+ response_media_type: Optional[str] = None,
+ status_code: Optional[int] = 200,
+ router_tags: Optional[List[str]] = None,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ if not endpoint.startswith("/"):
+ endpoint = "/" + endpoint
+ self._endpoint = endpoint
+ self._methods = methods
+ self._req_body = request_body
+ self._streaming_response = streaming_response
+ self._response_model = response_model
+ self._status_code = status_code
+ self._router_tags = router_tags
+ self._response_headers = response_headers
+ self._response_media_type = response_media_type
+ self._end_node: BaseOperator = None
+
+ async def trigger(self) -> None:
+ pass
+
+ def mount_to_router(self, router: "APIRouter") -> None:
+ from fastapi import Depends
+
+ methods = self._methods if isinstance(self._methods, list) else [self._methods]
+
+ def create_route_function(name, req_body_cls: Optional[Type[BaseModel]]):
+ async def _request_body_dependency(request: Request):
+ return await _parse_request_body(request, self._req_body)
+
+ async def route_function(body=Depends(_request_body_dependency)):
+ return await _trigger_dag(
+ body,
+ self.dag,
+ self._streaming_response,
+ self._response_headers,
+ self._response_media_type,
+ )
+
+ route_function.__name__ = name
+ return route_function
+
+ function_name = f"AWEL_trigger_route_{self._endpoint.replace('/', '_')}"
+ request_model = (
+ self._req_body
+ if isinstance(self._req_body, type)
+ and issubclass(self._req_body, BaseModel)
+ else None
+ )
+ dynamic_route_function = create_route_function(function_name, request_model)
+ logger.info(
+ f"mount router function {dynamic_route_function}({function_name}), endpoint: {self._endpoint}, methods: {methods}"
+ )
+
+ router.api_route(
+ self._endpoint,
+ methods=methods,
+ response_model=self._response_model,
+ status_code=self._status_code,
+ tags=self._router_tags,
+ )(dynamic_route_function)
+
+
+async def _parse_request_body(
+ request: Request, request_body_cls: Optional[Type[BaseModel]]
+):
+ if not request_body_cls:
+ return None
+ if request.method == "POST":
+ json_data = await request.json()
+ return request_body_cls(**json_data)
+ elif request.method == "GET":
+ return request_body_cls(**request.query_params)
+ else:
+ return request
+
+
+async def _trigger_dag(
+ body: Any,
+ dag: DAG,
+ streaming_response: Optional[bool] = False,
+ response_headers: Optional[Dict[str, str]] = None,
+ response_media_type: Optional[str] = None,
+) -> Any:
+ from fastapi.responses import StreamingResponse
+
+ end_node = dag.leaf_nodes
+ if len(end_node) != 1:
+ raise ValueError("HttpTrigger just support one leaf node in dag")
+ end_node = end_node[0]
+ if not streaming_response:
+ return await end_node.call(call_data={"data": body})
+ else:
+ headers = response_headers
+ media_type = response_media_type if response_media_type else "text/event-stream"
+ if not headers:
+ headers = {
+ "Content-Type": "text/event-stream",
+ "Cache-Control": "no-cache",
+ "Connection": "keep-alive",
+ "Transfer-Encoding": "chunked",
+ }
+ return StreamingResponse(
+ end_node.call_stream(call_data={"data": body}),
+ headers=headers,
+ media_type=media_type,
+ )
diff --git a/pilot/awel/trigger/trigger_manager.py b/pilot/awel/trigger/trigger_manager.py
new file mode 100644
index 000000000..feb674ffb
--- /dev/null
+++ b/pilot/awel/trigger/trigger_manager.py
@@ -0,0 +1,74 @@
+from abc import ABC, abstractmethod
+from typing import Any, TYPE_CHECKING, Optional
+import logging
+
+if TYPE_CHECKING:
+ from fastapi import APIRouter
+
+from pilot.component import SystemApp, BaseComponent, ComponentType
+
+logger = logging.getLogger(__name__)
+
+
+class TriggerManager(ABC):
+ @abstractmethod
+ def register_trigger(self, trigger: Any) -> None:
+ """ "Register a trigger to current manager"""
+
+
+class HttpTriggerManager(TriggerManager):
+ def __init__(
+ self,
+ router: Optional["APIRouter"] = None,
+ router_prefix: Optional[str] = "/api/v1/awel/trigger",
+ ) -> None:
+ if not router:
+ from fastapi import APIRouter
+
+ router = APIRouter()
+ self._router_prefix = router_prefix
+ self._router = router
+ self._trigger_map = {}
+
+ def register_trigger(self, trigger: Any) -> None:
+ from .http_trigger import HttpTrigger
+
+ if not isinstance(trigger, HttpTrigger):
+ raise ValueError(f"Current trigger {trigger} not an object of HttpTrigger")
+ trigger: HttpTrigger = trigger
+ trigger_id = trigger.node_id
+ if trigger_id not in self._trigger_map:
+ trigger.mount_to_router(self._router)
+ self._trigger_map[trigger_id] = trigger
+
+ def _init_app(self, system_app: SystemApp):
+ logger.info(
+ f"Include router {self._router} to prefix path {self._router_prefix}"
+ )
+ system_app.app.include_router(
+ self._router, prefix=self._router_prefix, tags=["AWEL"]
+ )
+
+
+class DefaultTriggerManager(TriggerManager, BaseComponent):
+ name = ComponentType.AWEL_TRIGGER_MANAGER
+
+ def __init__(self, system_app: SystemApp | None = None):
+ self.system_app = system_app
+ self.http_trigger = HttpTriggerManager()
+ super().__init__(None)
+
+ def init_app(self, system_app: SystemApp):
+ self.system_app = system_app
+
+ def register_trigger(self, trigger: Any) -> None:
+ from .http_trigger import HttpTrigger
+
+ if isinstance(trigger, HttpTrigger):
+ logger.info(f"Register trigger {trigger}")
+ self.http_trigger.register_trigger(trigger)
+ else:
+ raise ValueError(f"Unsupport trigger: {trigger}")
+
+ def after_register(self) -> None:
+ self.http_trigger._init_app(self.system_app)
diff --git a/pilot/base_modules/agent/commands/command_mange.py b/pilot/base_modules/agent/commands/command_mange.py
index be9e02811..89aa2f5f4 100644
--- a/pilot/base_modules/agent/commands/command_mange.py
+++ b/pilot/base_modules/agent/commands/command_mange.py
@@ -5,7 +5,9 @@ import time
import json
import logging
import xml.etree.ElementTree as ET
+import pandas as pd
+from pilot.common.json_utils import serialize
from datetime import datetime
from typing import Any, Callable, Optional, List
from pydantic import BaseModel
@@ -184,6 +186,8 @@ class PluginStatus(BaseModel):
start_time = datetime.now().timestamp() * 1000
end_time: int = None
+ df: Any = None
+
class ApiCall:
agent_prefix = "
[INST] <",
+ stop_token_ids=[2],
+ system_formatter=lambda msg: f"", "[UNK]"],
+ ),
+ override=True,
+)
+# source: https://huggingface.co/BAAI/AquilaChat2-34B/blob/4608b75855334b93329a771aee03869dbf7d88cc/predict.py#L227
+register_conv_template(
+ Conversation(
+ name="aquila",
+ system_message="A chat between a curious human and an artificial intelligence assistant. "
+ "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+ roles=("Human", "Assistant", "System"),
+ messages=(),
+ offset=0,
+ sep_style=SeparatorStyle.ADD_COLON_TWO,
+ sep="###",
+ sep2="",
+ stop_str=["", "[UNK]"],
+ ),
+ override=True,
+)
+# source: https://huggingface.co/BAAI/AquilaChat2-34B/blob/4608b75855334b93329a771aee03869dbf7d88cc/predict.py#L242
+register_conv_template(
+ Conversation(
+ name="aquila-v1",
+ roles=("<|startofpiece|>", "<|endofpiece|>", ""),
+ messages=(),
+ offset=0,
+ sep_style=SeparatorStyle.NO_COLON_TWO,
+ sep="",
+ sep2="",
+ stop_str=["", "<|endoftext|>"],
),
override=True,
)
diff --git a/pilot/model/operator/__init__.py b/pilot/model/operator/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/pilot/model/operator/model_operator.py b/pilot/model/operator/model_operator.py
new file mode 100644
index 000000000..d8ee62172
--- /dev/null
+++ b/pilot/model/operator/model_operator.py
@@ -0,0 +1,311 @@
+from typing import AsyncIterator, Dict, List, Union
+import logging
+from pilot.awel import (
+ BranchFunc,
+ StreamifyAbsOperator,
+ BranchOperator,
+ MapOperator,
+ TransformStreamAbsOperator,
+)
+from pilot.component import ComponentType
+from pilot.awel.operator.base import BaseOperator
+from pilot.model.base import ModelOutput
+from pilot.model.cluster import WorkerManager, WorkerManagerFactory
+from pilot.cache import LLMCacheClient, CacheManager, LLMCacheKey, LLMCacheValue
+
+logger = logging.getLogger(__name__)
+
+_LLM_MODEL_INPUT_VALUE_KEY = "llm_model_input_value"
+_LLM_MODEL_OUTPUT_CACHE_KEY = "llm_model_output_cache"
+
+
+class ModelStreamOperator(StreamifyAbsOperator[Dict, ModelOutput]):
+ """Operator for streaming processing of model outputs.
+
+ Args:
+ worker_manager (WorkerManager): The manager that handles worker processes for model inference.
+ **kwargs: Additional keyword arguments.
+
+ Methods:
+ streamify: Asynchronously processes a stream of inputs, yielding model outputs.
+ """
+
+ def __init__(self, worker_manager: WorkerManager = None, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.worker_manager = worker_manager
+
+ async def streamify(self, input_value: Dict) -> AsyncIterator[ModelOutput]:
+ """Process inputs as a stream and yield model outputs.
+
+ Args:
+ input_value (Dict): The input value for the model.
+
+ Returns:
+ AsyncIterator[ModelOutput]: An asynchronous iterator of model outputs.
+ """
+ if not self.worker_manager:
+ self.worker_manager = self.system_app.get_component(
+ ComponentType.WORKER_MANAGER_FACTORY, WorkerManagerFactory
+ ).create()
+ async for out in self.worker_manager.generate_stream(input_value):
+ yield out
+
+
+class ModelOperator(MapOperator[Dict, ModelOutput]):
+ """Operator for map-based processing of model outputs.
+
+ Args:
+ worker_manager (WorkerManager): Manager for handling worker processes.
+ **kwargs: Additional keyword arguments.
+
+ Methods:
+ map: Asynchronously processes a single input and returns the model output.
+ """
+
+ def __init__(self, worker_manager: WorkerManager = None, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.worker_manager = worker_manager
+
+ async def map(self, input_value: Dict) -> ModelOutput:
+ """Process a single input and return the model output.
+
+ Args:
+ input_value (Dict): The input value for the model.
+
+ Returns:
+ ModelOutput: The output from the model.
+ """
+ if not self.worker_manager:
+ self.worker_manager = self.system_app.get_component(
+ ComponentType.WORKER_MANAGER_FACTORY, WorkerManagerFactory
+ ).create()
+ return await self.worker_manager.generate(input_value)
+
+
+class CachedModelStreamOperator(StreamifyAbsOperator[Dict, ModelOutput]):
+ """Operator for streaming processing of model outputs with caching.
+
+ Args:
+ cache_manager (CacheManager): The cache manager to handle caching operations.
+ **kwargs: Additional keyword arguments.
+
+ Methods:
+ streamify: Processes a stream of inputs with cache support, yielding model outputs.
+ """
+
+ def __init__(self, cache_manager: CacheManager, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self._cache_manager = cache_manager
+ self._client = LLMCacheClient(cache_manager)
+
+ async def streamify(self, input_value: Dict) -> AsyncIterator[ModelOutput]:
+ """Process inputs as a stream with cache support and yield model outputs.
+
+ Args:
+ input_value (Dict): The input value for the model.
+
+ Returns:
+ AsyncIterator[ModelOutput]: An asynchronous iterator of model outputs.
+ """
+ cache_dict = _parse_cache_key_dict(input_value)
+ llm_cache_key: LLMCacheKey = self._client.new_key(**cache_dict)
+ llm_cache_value: LLMCacheValue = await self._client.get(llm_cache_key)
+ logger.info(f"llm_cache_value: {llm_cache_value}")
+ for out in llm_cache_value.get_value().output:
+ yield out
+
+
+class CachedModelOperator(MapOperator[Dict, ModelOutput]):
+ """Operator for map-based processing of model outputs with caching.
+
+ Args:
+ cache_manager (CacheManager): Manager for caching operations.
+ **kwargs: Additional keyword arguments.
+
+ Methods:
+ map: Processes a single input with cache support and returns the model output.
+ """
+
+ def __init__(self, cache_manager: CacheManager, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self._cache_manager = cache_manager
+ self._client = LLMCacheClient(cache_manager)
+
+ async def map(self, input_value: Dict) -> ModelOutput:
+ """Process a single input with cache support and return the model output.
+
+ Args:
+ input_value (Dict): The input value for the model.
+
+ Returns:
+ ModelOutput: The output from the model.
+ """
+ cache_dict = _parse_cache_key_dict(input_value)
+ llm_cache_key: LLMCacheKey = self._client.new_key(**cache_dict)
+ llm_cache_value: LLMCacheValue = await self._client.get(llm_cache_key)
+ logger.info(f"llm_cache_value: {llm_cache_value}")
+ return llm_cache_value.get_value().output
+
+
+class ModelCacheBranchOperator(BranchOperator[Dict, Dict]):
+ """
+ A branch operator that decides whether to use cached data or to process data using the model.
+
+ Args:
+ cache_manager (CacheManager): The cache manager for managing cache operations.
+ model_task_name (str): The name of the task to process data using the model.
+ cache_task_name (str): The name of the task to process data using the cache.
+ **kwargs: Additional keyword arguments.
+ """
+
+ def __init__(
+ self,
+ cache_manager: CacheManager,
+ model_task_name: str,
+ cache_task_name: str,
+ **kwargs,
+ ):
+ super().__init__(branches=None, **kwargs)
+ self._cache_manager = cache_manager
+ self._client = LLMCacheClient(cache_manager)
+ self._model_task_name = model_task_name
+ self._cache_task_name = cache_task_name
+
+ async def branchs(self) -> Dict[BranchFunc[Dict], Union[BaseOperator, str]]:
+ """Defines branch logic based on cache availability.
+
+ Returns:
+ Dict[BranchFunc[Dict], Union[BaseOperator, str]]: A dictionary mapping branch functions to task names.
+ """
+
+ async def check_cache_true(input_value: Dict) -> bool:
+ # Check if the cache contains the result for the given input
+ if not input_value["model_cache_enable"]:
+ return False
+ cache_dict = _parse_cache_key_dict(input_value)
+ cache_key: LLMCacheKey = self._client.new_key(**cache_dict)
+ cache_value = await self._client.get(cache_key)
+ logger.debug(
+ f"cache_key: {cache_key}, hash key: {hash(cache_key)}, cache_value: {cache_value}"
+ )
+ await self.current_dag_context.save_to_share_data(
+ _LLM_MODEL_INPUT_VALUE_KEY, cache_key
+ )
+ return True if cache_value else False
+
+ async def check_cache_false(input_value: Dict):
+ # Inverse of check_cache_true
+ return not await check_cache_true(input_value)
+
+ return {
+ check_cache_true: self._cache_task_name,
+ check_cache_false: self._model_task_name,
+ }
+
+
+class ModelStreamSaveCacheOperator(
+ TransformStreamAbsOperator[ModelOutput, ModelOutput]
+):
+ """An operator to save the stream of model outputs to cache.
+
+ Args:
+ cache_manager (CacheManager): The cache manager for handling cache operations.
+ **kwargs: Additional keyword arguments.
+ """
+
+ def __init__(self, cache_manager: CacheManager, **kwargs):
+ self._cache_manager = cache_manager
+ self._client = LLMCacheClient(cache_manager)
+ super().__init__(**kwargs)
+
+ async def transform_stream(
+ self, input_value: AsyncIterator[ModelOutput]
+ ) -> AsyncIterator[ModelOutput]:
+ """Transforms the input stream by saving the outputs to cache.
+
+ Args:
+ input_value (AsyncIterator[ModelOutput]): An asynchronous iterator of model outputs.
+
+ Returns:
+ AsyncIterator[ModelOutput]: The same input iterator, but the outputs are saved to cache.
+ """
+ llm_cache_key: LLMCacheKey = None
+ outputs = []
+ async for out in input_value:
+ if not llm_cache_key:
+ llm_cache_key = await self.current_dag_context.get_share_data(
+ _LLM_MODEL_INPUT_VALUE_KEY
+ )
+ outputs.append(out)
+ yield out
+ if llm_cache_key and _is_success_model_output(outputs):
+ llm_cache_value: LLMCacheValue = self._client.new_value(output=outputs)
+ await self._client.set(llm_cache_key, llm_cache_value)
+
+
+class ModelSaveCacheOperator(MapOperator[ModelOutput, ModelOutput]):
+ """An operator to save a single model output to cache.
+
+ Args:
+ cache_manager (CacheManager): The cache manager for handling cache operations.
+ **kwargs: Additional keyword arguments.
+ """
+
+ def __init__(self, cache_manager: CacheManager, **kwargs):
+ self._cache_manager = cache_manager
+ self._client = LLMCacheClient(cache_manager)
+ super().__init__(**kwargs)
+
+ async def map(self, input_value: ModelOutput) -> ModelOutput:
+ """Saves a single model output to cache and returns it.
+
+ Args:
+ input_value (ModelOutput): The output from the model to be cached.
+
+ Returns:
+ ModelOutput: The same input model output.
+ """
+ llm_cache_key: LLMCacheKey = await self.current_dag_context.get_share_data(
+ _LLM_MODEL_INPUT_VALUE_KEY
+ )
+ llm_cache_value: LLMCacheValue = self._client.new_value(output=input_value)
+ if llm_cache_key and _is_success_model_output(input_value):
+ await self._client.set(llm_cache_key, llm_cache_value)
+ return input_value
+
+
+def _parse_cache_key_dict(input_value: Dict) -> Dict:
+ """Parses and extracts relevant fields from input to form a cache key dictionary.
+
+ Args:
+ input_value (Dict): The input dictionary containing model and prompt parameters.
+
+ Returns:
+ Dict: A dictionary used for generating cache keys.
+ """
+ prompt: str = input_value.get("prompt")
+ if prompt:
+ prompt = prompt.strip()
+ return {
+ "prompt": prompt,
+ "model_name": input_value.get("model"),
+ "temperature": input_value.get("temperature"),
+ "max_new_tokens": input_value.get("max_new_tokens"),
+ "top_p": input_value.get("top_p", "1.0"),
+ # TODO pass model_type
+ "model_type": input_value.get("model_type", "huggingface"),
+ }
+
+
+def _is_success_model_output(out: Union[Dict, ModelOutput, List[ModelOutput]]) -> bool:
+ if not out:
+ return False
+ if isinstance(out, list):
+ # check last model output
+ out = out[-1]
+ error_code = 0
+ if isinstance(out, ModelOutput):
+ error_code = out.error_code
+ else:
+ error_code = int(out.get("error_code", 0))
+ return error_code == 0
diff --git a/pilot/model/parameter.py b/pilot/model/parameter.py
index ba0000435..79558e02c 100644
--- a/pilot/model/parameter.py
+++ b/pilot/model/parameter.py
@@ -1,9 +1,10 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
+
import os
from dataclasses import dataclass, field
from enum import Enum
-from typing import Dict, Optional
+from typing import Dict, Optional, Union, Tuple
from pilot.model.conversation import conv_templates
from pilot.utils.parameter_utils import BaseParameters
@@ -19,6 +20,35 @@ class WorkerType(str, Enum):
def values():
return [item.value for item in WorkerType]
+ @staticmethod
+ def to_worker_key(worker_name, worker_type: Union[str, "WorkerType"]) -> str:
+ """Generate worker key from worker name and worker type
+
+ Args:
+ worker_name (str): Worker name(eg., chatglm2-6b)
+ worker_type (Union[str, "WorkerType"]): Worker type(eg., 'llm', or [`WorkerType.LLM`])
+
+ Returns:
+ str: Generated worker key
+ """
+ if "@" in worker_name:
+ raise ValueError(f"Invaild symbol '@' in your worker name {worker_name}")
+ if isinstance(worker_type, WorkerType):
+ worker_type = worker_type.value
+ return f"{worker_name}@{worker_type}"
+
+ @staticmethod
+ def parse_worker_key(worker_key: str) -> Tuple[str, str]:
+ """Parse worker name and worker type from worker key
+
+ Args:
+ worker_key (str): Worker key generated by [`WorkerType.to_worker_key`]
+
+ Returns:
+ Tuple[str, str]: Worker name and worker type
+ """
+ return tuple(worker_key.split("@"))
+
@dataclass
class ModelControllerParameters(BaseParameters):
@@ -58,6 +88,68 @@ class ModelControllerParameters(BaseParameters):
"help": "The filename to store tracer span records",
},
)
+ tracer_storage_cls: Optional[str] = field(
+ default=None,
+ metadata={
+ "help": "The storage class to storage tracer span records",
+ },
+ )
+
+
+@dataclass
+class ModelAPIServerParameters(BaseParameters):
+ host: Optional[str] = field(
+ default="0.0.0.0", metadata={"help": "Model API server deploy host"}
+ )
+ port: Optional[int] = field(
+ default=8100, metadata={"help": "Model API server deploy port"}
+ )
+ daemon: Optional[bool] = field(
+ default=False, metadata={"help": "Run Model API server in background"}
+ )
+ controller_addr: Optional[str] = field(
+ default="http://127.0.0.1:8000",
+ metadata={"help": "The Model controller address to connect"},
+ )
+
+ api_keys: Optional[str] = field(
+ default=None,
+ metadata={"help": "Optional list of comma separated API keys"},
+ )
+
+ log_level: Optional[str] = field(
+ default=None,
+ metadata={
+ "help": "Logging level",
+ "valid_values": [
+ "FATAL",
+ "ERROR",
+ "WARNING",
+ "WARNING",
+ "INFO",
+ "DEBUG",
+ "NOTSET",
+ ],
+ },
+ )
+ log_file: Optional[str] = field(
+ default="dbgpt_model_apiserver.log",
+ metadata={
+ "help": "The filename to store log",
+ },
+ )
+ tracer_file: Optional[str] = field(
+ default="dbgpt_model_apiserver_tracer.jsonl",
+ metadata={
+ "help": "The filename to store tracer span records",
+ },
+ )
+ tracer_storage_cls: Optional[str] = field(
+ default=None,
+ metadata={
+ "help": "The storage class to storage tracer span records",
+ },
+ )
@dataclass
@@ -146,6 +238,12 @@ class ModelWorkerParameters(BaseModelParameters):
"help": "The filename to store tracer span records",
},
)
+ tracer_storage_cls: Optional[str] = field(
+ default=None,
+ metadata={
+ "help": "The storage class to storage tracer span records",
+ },
+ )
@dataclass
@@ -318,6 +416,13 @@ class ProxyModelParameters(BaseModelParameters):
},
)
+ proxy_api_app_id: Optional[str] = field(
+ default=None,
+ metadata={
+ "help": "The app id for current proxy LLM(Just for spark proxy LLM now)."
+ },
+ )
+
proxy_api_type: Optional[str] = field(
default=None,
metadata={
diff --git a/pilot/model/proxy/llms/chatgpt.py b/pilot/model/proxy/llms/chatgpt.py
index 994986b5c..4c2a371ce 100644
--- a/pilot/model/proxy/llms/chatgpt.py
+++ b/pilot/model/proxy/llms/chatgpt.py
@@ -4,18 +4,25 @@
import os
from typing import List
import logging
-
-import openai
-
+import importlib.metadata as metadata
from pilot.model.proxy.llms.proxy_model import ProxyModel
from pilot.model.parameter import ProxyModelParameters
from pilot.scene.base_message import ModelMessage, ModelMessageRoleType
from pilot.server.monitor.api_key_db import ApiKeyDao
+import httpx
logger = logging.getLogger(__name__)
def _initialize_openai(params: ProxyModelParameters):
+ try:
+ import openai
+ except ImportError as exc:
+ raise ValueError(
+ "Could not import python package: openai "
+ "Please install openai by command `pip install openai` "
+ ) from exc
+
api_type = params.proxy_api_type or os.getenv("OPENAI_API_TYPE", "open_ai")
api_base = params.proxy_api_base or os.getenv(
@@ -55,14 +62,45 @@ def _initialize_openai(params: ProxyModelParameters):
return openai_params
+def _initialize_openai_v1(params: ProxyModelParameters):
+ try:
+ from openai import OpenAI
+ except ImportError as exc:
+ raise ValueError(
+ "Could not import python package: openai "
+ "Please install openai by command `pip install openai"
+ )
+
+ api_type = params.proxy_api_type or os.getenv("OPENAI_API_TYPE", "open_ai")
+
+ base_url = params.proxy_api_base or os.getenv(
+ "OPENAI_API_TYPE",
+ os.getenv("AZURE_OPENAI_ENDPOINT") if api_type == "azure" else None,
+ )
+ api_key = params.proxy_api_key or os.getenv(
+ "OPENAI_API_KEY",
+ os.getenv("AZURE_OPENAI_KEY") if api_type == "azure" else None,
+ )
+ api_version = params.proxy_api_version or os.getenv("OPENAI_API_VERSION")
+
+ if not base_url and params.proxy_server_url:
+ # Adapt previous proxy_server_url configuration
+ base_url = params.proxy_server_url.split("/chat/completions")[0]
+
+ proxies = params.http_proxy
+ openai_params = {
+ "api_key": api_key,
+ "base_url": base_url,
+ }
+ return openai_params, api_type, api_version, proxies
+
+
def _build_request(model: ProxyModel, params):
history = []
model_params = model.get_params()
logger.info(f"Model: {model}, model_params: {model_params}")
- openai_params = _initialize_openai(model_params)
-
messages: List[ModelMessage] = params["messages"]
# Add history conversation
for message in messages:
@@ -93,45 +131,113 @@ def _build_request(model: ProxyModel, params):
}
proxyllm_backend = model_params.proxyllm_backend
- if openai_params["api_type"] == "azure":
- # engine = "deployment_name".
- proxyllm_backend = proxyllm_backend or "gpt-35-turbo"
- payloads["engine"] = proxyllm_backend
- else:
+ if metadata.version("openai") >= "1.0.0":
+ openai_params, api_type, api_version, proxies = _initialize_openai_v1(
+ model_params
+ )
proxyllm_backend = proxyllm_backend or "gpt-3.5-turbo"
payloads["model"] = proxyllm_backend
+ else:
+ openai_params = _initialize_openai(model_params)
+ if openai_params["api_type"] == "azure":
+ # engine = "deployment_name".
+ proxyllm_backend = proxyllm_backend or "gpt-35-turbo"
+ payloads["engine"] = proxyllm_backend
+ else:
+ proxyllm_backend = proxyllm_backend or "gpt-3.5-turbo"
+ payloads["model"] = proxyllm_backend
- logger.info(
- f"Send request to real model {proxyllm_backend}, openai_params: {openai_params}"
- )
+ logger.info(f"Send request to real model {proxyllm_backend}")
return history, payloads
def chatgpt_generate_stream(
model: ProxyModel, tokenizer, params, device, context_len=2048
):
- history, payloads = _build_request(model, params)
+ if metadata.version("openai") >= "1.0.0":
+ model_params = model.get_params()
+ openai_params, api_type, api_version, proxies = _initialize_openai_v1(
+ model_params
+ )
+ history, payloads = _build_request(model, params)
+ if api_type == "azure":
+ from openai import AzureOpenAI
- res = openai.ChatCompletion.create(messages=history, **payloads)
+ client = AzureOpenAI(
+ api_key=openai_params["api_key"],
+ api_version=api_version,
+ azure_endpoint=openai_params["base_url"],
+ http_client=httpx.Client(proxies=proxies),
+ )
+ else:
+ from openai import OpenAI
- text = ""
- for r in res:
- if r["choices"][0]["delta"].get("content") is not None:
- content = r["choices"][0]["delta"]["content"]
- text += content
- yield text
+ client = OpenAI(**openai_params, http_client=httpx.Client(proxies=proxies))
+ res = client.chat.completions.create(messages=history, **payloads)
+ text = ""
+ for r in res:
+ if r.choices[0].delta.content is not None:
+ content = r.choices[0].delta.content
+ text += content
+ yield text
+
+ else:
+ import openai
+
+ history, payloads = _build_request(model, params)
+
+ res = openai.ChatCompletion.create(messages=history, **payloads)
+
+ text = ""
+ for r in res:
+ if r["choices"][0]["delta"].get("content") is not None:
+ content = r["choices"][0]["delta"]["content"]
+ text += content
+ yield text
async def async_chatgpt_generate_stream(
model: ProxyModel, tokenizer, params, device, context_len=2048
):
- history, payloads = _build_request(model, params)
+ if metadata.version("openai") >= "1.0.0":
+ model_params = model.get_params()
+ openai_params, api_type, api_version, proxies = _initialize_openai_v1(
+ model_params
+ )
+ history, payloads = _build_request(model, params)
+ if api_type == "azure":
+ from openai import AsyncAzureOpenAI
- res = await openai.ChatCompletion.acreate(messages=history, **payloads)
+ client = AsyncAzureOpenAI(
+ api_key=openai_params["api_key"],
+ api_version=api_version,
+ azure_endpoint=openai_params["base_url"],
+ http_client=httpx.AsyncClient(proxies=proxies),
+ )
+ else:
+ from openai import AsyncOpenAI
- text = ""
- async for r in res:
- if r["choices"][0]["delta"].get("content") is not None:
- content = r["choices"][0]["delta"]["content"]
- text += content
- yield text
+ client = AsyncOpenAI(
+ **openai_params, http_client=httpx.AsyncClient(proxies=proxies)
+ )
+
+ res = await client.chat.completions.create(messages=history, **payloads)
+ text = ""
+ for r in res:
+ if r.choices[0].delta.content is not None:
+ content = r.choices[0].delta.content
+ text += content
+ yield text
+ else:
+ import openai
+
+ history, payloads = _build_request(model, params)
+
+ res = await openai.ChatCompletion.acreate(messages=history, **payloads)
+
+ text = ""
+ async for r in res:
+ if r["choices"][0]["delta"].get("content") is not None:
+ content = r["choices"][0]["delta"]["content"]
+ text += content
+ yield text
diff --git a/pilot/model/proxy/llms/tongyi.py b/pilot/model/proxy/llms/tongyi.py
index fb826e49c..5bc0a97f3 100644
--- a/pilot/model/proxy/llms/tongyi.py
+++ b/pilot/model/proxy/llms/tongyi.py
@@ -7,6 +7,35 @@ from pilot.scene.base_message import ModelMessage, ModelMessageRoleType
logger = logging.getLogger(__name__)
+def __convert_2_tongyi_messages(messages: List[ModelMessage]):
+ chat_round = 0
+ tongyi_messages = []
+
+ last_usr_message = ""
+ system_messages = []
+
+ for message in messages:
+ if message.role == ModelMessageRoleType.HUMAN:
+ last_usr_message = message.content
+ elif message.role == ModelMessageRoleType.SYSTEM:
+ system_messages.append(message.content)
+ elif message.role == ModelMessageRoleType.AI:
+ last_ai_message = message.content
+ tongyi_messages.append({"role": "user", "content": last_usr_message})
+ tongyi_messages.append({"role": "assistant", "content": last_ai_message})
+ if len(system_messages) > 0:
+ if len(system_messages) < 2:
+ tongyi_messages.insert(0, {"role": "system", "content": system_messages[0]})
+ else:
+ tongyi_messages.append({"role": "user", "content": system_messages[1]})
+ else:
+ last_message = messages[-1]
+ if last_message.role == ModelMessageRoleType.HUMAN:
+ tongyi_messages.append({"role": "user", "content": last_message.content})
+
+ return tongyi_messages
+
+
def tongyi_generate_stream(
model: ProxyModel, tokenizer, params, device, context_len=2048
):
@@ -23,40 +52,9 @@ def tongyi_generate_stream(
if not proxyllm_backend:
proxyllm_backend = Generation.Models.qwen_turbo # By Default qwen_turbo
- history = []
-
messages: List[ModelMessage] = params["messages"]
- # Add history conversation
-
- if len(messages) > 1 and messages[0].role == ModelMessageRoleType.SYSTEM:
- role_define = messages.pop(0)
- history.append({"role": "system", "content": role_define.content})
- else:
- message = messages.pop(0)
- if message.role == ModelMessageRoleType.HUMAN:
- history.append({"role": "user", "content": message.content})
- for message in messages:
- if message.role == ModelMessageRoleType.SYSTEM:
- history.append({"role": "user", "content": message.content})
- # elif message.role == ModelMessageRoleType.HUMAN:
- # history.append({"role": "user", "content": message.content})
- elif message.role == ModelMessageRoleType.AI:
- history.append({"role": "assistant", "content": message.content})
- else:
- pass
-
- # temp_his = history[::-1]
- temp_his = history
- last_user_input = None
- for m in temp_his:
- if m["role"] == "user":
- last_user_input = m
- break
-
- if last_user_input:
- history.remove(last_user_input)
- history.append(last_user_input)
+ history = __convert_2_tongyi_messages(messages)
gen = Generation()
res = gen.call(
proxyllm_backend,
diff --git a/pilot/model/proxy/llms/wenxin.py b/pilot/model/proxy/llms/wenxin.py
index acc82907c..cfd47fd18 100644
--- a/pilot/model/proxy/llms/wenxin.py
+++ b/pilot/model/proxy/llms/wenxin.py
@@ -26,6 +26,41 @@ def _build_access_token(api_key: str, secret_key: str) -> str:
return res.json().get("access_token")
+def __convert_2_wenxin_messages(messages: List[ModelMessage]):
+ chat_round = 0
+ wenxin_messages = []
+
+ last_usr_message = ""
+ system_messages = []
+
+ for message in messages:
+ if message.role == ModelMessageRoleType.HUMAN:
+ last_usr_message = message.content
+ elif message.role == ModelMessageRoleType.SYSTEM:
+ system_messages.append(message.content)
+ elif message.role == ModelMessageRoleType.AI:
+ last_ai_message = message.content
+ wenxin_messages.append({"role": "user", "content": last_usr_message})
+ wenxin_messages.append({"role": "assistant", "content": last_ai_message})
+
+ # build last user messge
+
+ if len(system_messages) > 0:
+ if len(system_messages) > 1:
+ end_message = system_messages[-1]
+ else:
+ last_message = messages[-1]
+ if last_message.role == ModelMessageRoleType.HUMAN:
+ end_message = system_messages[-1] + "\n" + last_message.content
+ else:
+ end_message = system_messages[-1]
+ else:
+ last_message = messages[-1]
+ end_message = last_message.content
+ wenxin_messages.append({"role": "user", "content": end_message})
+ return wenxin_messages, system_messages
+
+
def wenxin_generate_stream(
model: ProxyModel, tokenizer, params, device, context_len=2048
):
@@ -40,8 +75,9 @@ def wenxin_generate_stream(
if not model_version:
yield f"Unsupport model version {model_name}"
- proxy_api_key = model_params.proxy_api_key
- proxy_api_secret = model_params.proxy_api_secret
+ keys: [] = model_params.proxy_api_key.split(";")
+ proxy_api_key = keys[0]
+ proxy_api_secret = keys[1]
access_token = _build_access_token(proxy_api_key, proxy_api_secret)
headers = {"Content-Type": "application/json", "Accept": "application/json"}
@@ -51,40 +87,42 @@ def wenxin_generate_stream(
if not access_token:
yield "Failed to get access token. please set the correct api_key and secret key."
- history = []
-
messages: List[ModelMessage] = params["messages"]
# Add history conversation
+ # system = ""
+ # if len(messages) > 1 and messages[0].role == ModelMessageRoleType.SYSTEM:
+ # role_define = messages.pop(0)
+ # system = role_define.content
+ # else:
+ # message = messages.pop(0)
+ # if message.role == ModelMessageRoleType.HUMAN:
+ # history.append({"role": "user", "content": message.content})
+ # for message in messages:
+ # if message.role == ModelMessageRoleType.SYSTEM:
+ # history.append({"role": "user", "content": message.content})
+ # # elif message.role == ModelMessageRoleType.HUMAN:
+ # # history.append({"role": "user", "content": message.content})
+ # elif message.role == ModelMessageRoleType.AI:
+ # history.append({"role": "assistant", "content": message.content})
+ # else:
+ # pass
+ #
+ # # temp_his = history[::-1]
+ # temp_his = history
+ # last_user_input = None
+ # for m in temp_his:
+ # if m["role"] == "user":
+ # last_user_input = m
+ # break
+ #
+ # if last_user_input:
+ # history.remove(last_user_input)
+ # history.append(last_user_input)
+ #
+ history, systems = __convert_2_wenxin_messages(messages)
system = ""
- if len(messages) > 1 and messages[0].role == ModelMessageRoleType.SYSTEM:
- role_define = messages.pop(0)
- system = role_define.content
- else:
- message = messages.pop(0)
- if message.role == ModelMessageRoleType.HUMAN:
- history.append({"role": "user", "content": message.content})
- for message in messages:
- if message.role == ModelMessageRoleType.SYSTEM:
- history.append({"role": "user", "content": message.content})
- # elif message.role == ModelMessageRoleType.HUMAN:
- # history.append({"role": "user", "content": message.content})
- elif message.role == ModelMessageRoleType.AI:
- history.append({"role": "assistant", "content": message.content})
- else:
- pass
-
- # temp_his = history[::-1]
- temp_his = history
- last_user_input = None
- for m in temp_his:
- if m["role"] == "user":
- last_user_input = m
- break
-
- if last_user_input:
- history.remove(last_user_input)
- history.append(last_user_input)
-
+ if systems and len(systems) > 0:
+ system = systems[0]
payload = {
"messages": history,
"system": system,
diff --git a/pilot/model/proxy/llms/zhipu.py b/pilot/model/proxy/llms/zhipu.py
index 89e7dd9a0..c5fabe8ed 100644
--- a/pilot/model/proxy/llms/zhipu.py
+++ b/pilot/model/proxy/llms/zhipu.py
@@ -8,6 +8,41 @@ from pilot.scene.base_message import ModelMessage, ModelMessageRoleType
CHATGLM_DEFAULT_MODEL = "chatglm_pro"
+def __convert_2_wenxin_messages(messages: List[ModelMessage]):
+ chat_round = 0
+ wenxin_messages = []
+
+ last_usr_message = ""
+ system_messages = []
+
+ for message in messages:
+ if message.role == ModelMessageRoleType.HUMAN:
+ last_usr_message = message.content
+ elif message.role == ModelMessageRoleType.SYSTEM:
+ system_messages.append(message.content)
+ elif message.role == ModelMessageRoleType.AI:
+ last_ai_message = message.content
+ wenxin_messages.append({"role": "user", "content": last_usr_message})
+ wenxin_messages.append({"role": "assistant", "content": last_ai_message})
+
+ # build last user messge
+
+ if len(system_messages) > 0:
+ if len(system_messages) > 1:
+ end_message = system_messages[-1]
+ else:
+ last_message = messages[-1]
+ if last_message.role == ModelMessageRoleType.HUMAN:
+ end_message = system_messages[-1] + "\n" + last_message.content
+ else:
+ end_message = system_messages[-1]
+ else:
+ last_message = messages[-1]
+ end_message = last_message.content
+ wenxin_messages.append({"role": "user", "content": end_message})
+ return wenxin_messages, system_messages
+
+
def zhipu_generate_stream(
model: ProxyModel, tokenizer, params, device, context_len=2048
):
@@ -22,40 +57,40 @@ def zhipu_generate_stream(
import zhipuai
zhipuai.api_key = proxy_api_key
- history = []
messages: List[ModelMessage] = params["messages"]
# Add history conversation
- system = ""
- if len(messages) > 1 and messages[0].role == ModelMessageRoleType.SYSTEM:
- role_define = messages.pop(0)
- system = role_define.content
- else:
- message = messages.pop(0)
- if message.role == ModelMessageRoleType.HUMAN:
- history.append({"role": "user", "content": message.content})
- for message in messages:
- if message.role == ModelMessageRoleType.SYSTEM:
- history.append({"role": "user", "content": message.content})
- # elif message.role == ModelMessageRoleType.HUMAN:
- # history.append({"role": "user", "content": message.content})
- elif message.role == ModelMessageRoleType.AI:
- history.append({"role": "assistant", "content": message.content})
- else:
- pass
-
- # temp_his = history[::-1]
- temp_his = history
- last_user_input = None
- for m in temp_his:
- if m["role"] == "user":
- last_user_input = m
- break
-
- if last_user_input:
- history.remove(last_user_input)
- history.append(last_user_input)
+ # system = ""
+ # if len(messages) > 1 and messages[0].role == ModelMessageRoleType.SYSTEM:
+ # role_define = messages.pop(0)
+ # system = role_define.content
+ # else:
+ # message = messages.pop(0)
+ # if message.role == ModelMessageRoleType.HUMAN:
+ # history.append({"role": "user", "content": message.content})
+ # for message in messages:
+ # if message.role == ModelMessageRoleType.SYSTEM:
+ # history.append({"role": "user", "content": message.content})
+ # # elif message.role == ModelMessageRoleType.HUMAN:
+ # # history.append({"role": "user", "content": message.content})
+ # elif message.role == ModelMessageRoleType.AI:
+ # history.append({"role": "assistant", "content": message.content})
+ # else:
+ # pass
+ #
+ # # temp_his = history[::-1]
+ # temp_his = history
+ # last_user_input = None
+ # for m in temp_his:
+ # if m["role"] == "user":
+ # last_user_input = m
+ # break
+ #
+ # if last_user_input:
+ # history.remove(last_user_input)
+ # history.append(last_user_input)
+ history, systems = __convert_2_wenxin_messages(messages)
res = zhipuai.model_api.sse_invoke(
model=proxyllm_backend,
prompt=history,
diff --git a/pilot/openapi/api_v1/api_v1.py b/pilot/openapi/api_v1/api_v1.py
index 640033ea1..fde191953 100644
--- a/pilot/openapi/api_v1/api_v1.py
+++ b/pilot/openapi/api_v1/api_v1.py
@@ -21,6 +21,7 @@ from fastapi.responses import StreamingResponse
from fastapi.exceptions import RequestValidationError
from typing import List
import tempfile
+from concurrent.futures import Executor
from pilot.component import ComponentType
from pilot.openapi.api_view_model import (
@@ -52,6 +53,8 @@ from pilot.memory.chat_history.chat_hisotry_factory import ChatHistory
from pilot.model.cluster import BaseModelController, WorkerManager, WorkerManagerFactory
from pilot.model.base import FlatSupportedModel
from pilot.user import UserDao, UserRequest, get_user_from_headers
+from pilot.utils.tracer import root_tracer, SpanType
+from pilot.utils.executor_utils import ExecutorFactory, blocking_func_to_async
router = APIRouter()
CFG = Config()
@@ -81,10 +84,13 @@ def __new_conversation(chat_mode, user_id) -> ConversationVo:
def get_db_list(user_id: str = None):
dbs = CFG.LOCAL_DB_MANAGE.get_db_list(user_id=user_id)
- params: dict = {}
+ db_params = []
for item in dbs:
- params.update({item["db_name"]: item["db_name"]})
- return params
+ params: dict = {}
+ params.update({"param": item["db_name"]})
+ params.update({"type": item["db_type"]})
+ db_params.append(params)
+ return db_params
def plugins_select_info():
@@ -116,12 +122,15 @@ def knowledge_list_info():
def knowledge_list():
"""return knowledge space list"""
- params: dict = {}
request = KnowledgeSpaceRequest()
spaces = knowledge_service.get_knowledge_space(request)
+ space_list = []
for space in spaces:
- params.update({space.name: space.name})
- return params
+ params: dict = {}
+ params.update({"param": space.name})
+ params.update({"type": "space"})
+ space_list.append(params)
+ return space_list
def get_model_controller() -> BaseModelController:
@@ -138,6 +147,13 @@ def get_worker_manager() -> WorkerManager:
return worker_manager
+def get_executor() -> Executor:
+ """Get the global default executor"""
+ return CFG.SYSTEM_APP.get_component(
+ ComponentType.EXECUTOR_DEFAULT, ExecutorFactory
+ ).create()
+
+
@router.get("/v1/chat/db/list", response_model=Result[DBConfig])
async def db_connect_list(user_token: UserRequest = Depends(get_user_from_headers)):
results = CFG.LOCAL_DB_MANAGE.get_db_list(user_token.user_id)
@@ -170,14 +186,16 @@ async def async_db_summary_embedding(db_name, db_type):
@router.post("/v1/chat/db/test/connect", response_model=Result[bool])
async def test_connect(db_config: DBConfig = Body()):
try:
+ # TODO Change the synchronous call to the asynchronous call
CFG.LOCAL_DB_MANAGE.test_connect(db_config)
return Result.succ(True)
except Exception as e:
- return Result.faild(code="E1001", msg=str(e))
+ return Result.failed(code="E1001", msg=str(e))
@router.post("/v1/chat/db/summary", response_model=Result[bool])
async def db_summary(db_name: str, db_type: str):
+ # TODO Change the synchronous call to the asynchronous call
async_db_summary_embedding(db_name, db_type)
return Result.succ(True)
@@ -227,8 +245,8 @@ async def dialogue_scenes():
scene_vos: List[ChatSceneVo] = []
new_modes: List[ChatScene] = [
ChatScene.ChatWithDbExecute,
- ChatScene.ChatExcel,
ChatScene.ChatWithDbQA,
+ ChatScene.ChatExcel,
ChatScene.ChatKnowledge,
ChatScene.ChatDashboard,
ChatScene.ChatAgent,
@@ -267,6 +285,10 @@ async def params_list(chat_mode: str = ChatScene.ChatNormal.value(), user_token:
result = plugins_select_info()
elif ChatScene.ChatKnowledge.value() == chat_mode:
result = knowledge_list()
+ elif ChatScene.ChatKnowledge.ExtractRefineSummary.value() == chat_mode:
+ result = knowledge_list()
+ else:
+ return Result.succ(None)
if result and result.get("dbgpt"):
del result["dbgpt"]
if result and result.get("auth"):
@@ -302,14 +324,14 @@ async def params_load(
select_param=doc_file.filename,
model_name=model_name,
)
- chat: BaseChat = get_chat_instance(dialogue, user_token.user_id)
+ chat: BaseChat = await get_chat_instance(dialogue, user_token.user_id)
resp = await chat.prepare()
### refresh messages
return Result.succ(get_hist_messages(conv_uid))
except Exception as e:
logger.error("excel load error!", e)
- return Result.faild(code="E000X", msg=f"File Load Error {e}")
+ return Result.failed(code="E000X", msg=f"File Load Error {e}")
@router.post("/v1/user/add")
@@ -330,6 +352,7 @@ async def add_user(user_req: UserRequest):
async def dialogue_delete(con_uid: str):
history_fac = ChatHistory()
history_mem = history_fac.get_store_instance(con_uid)
+ # TODO Change the synchronous call to the asynchronous call
history_mem.delete()
return Result.succ(None)
@@ -342,7 +365,6 @@ def get_hist_messages(conv_uid: str):
history_messages: List[OnceConversation] = history_mem.get_messages()
if history_messages:
for once in history_messages:
- print(f"once:{once}")
model_name = once.get("model_name", CFG.LLM_MODEL)
once_message_vos = [
message2Vo(element, once["chat_order"], model_name)
@@ -355,10 +377,12 @@ def get_hist_messages(conv_uid: str):
@router.get("/v1/chat/dialogue/messages/history", response_model=Result[MessageVo])
async def dialogue_history_messages(con_uid: str):
print(f"dialogue_history_messages:{con_uid}")
+ # TODO Change the synchronous call to the asynchronous call
return Result.succ(get_hist_messages(con_uid))
-def get_chat_instance(dialogue: ConversationVo = Body(), user_id: str = None) -> BaseChat:
+# def get_chat_instance(dialogue: ConversationVo = Body(), user_id: str = None) -> BaseChat:
+async def get_chat_instance(dialogue: ConversationVo = Body(), user_id: str = None) -> BaseChat:
logger.info(f"get_chat_instance:{dialogue}")
if not dialogue.chat_mode:
dialogue.chat_mode = ChatScene.ChatNormal.value()
@@ -368,7 +392,7 @@ def get_chat_instance(dialogue: ConversationVo = Body(), user_id: str = None) ->
if not ChatScene.is_valid_mode(dialogue.chat_mode):
raise StopAsyncIteration(
- Result.faild("Unsupported Chat Mode," + dialogue.chat_mode + "!")
+ Result.failed("Unsupported Chat Mode," + dialogue.chat_mode + "!")
)
chat_param = {
@@ -378,8 +402,14 @@ def get_chat_instance(dialogue: ConversationVo = Body(), user_id: str = None) ->
"model_name": dialogue.model_name,
"user_id": user_id,
}
- chat: BaseChat = CHAT_FACTORY.get_implementation(
- dialogue.chat_mode, **{"chat_param": chat_param}
+ # chat: BaseChat = CHAT_FACTORY.get_implementation(
+ # dialogue.chat_mode, **{"chat_param": chat_param}
+ # )
+ chat: BaseChat = await blocking_func_to_async(
+ get_executor(),
+ CHAT_FACTORY.get_implementation,
+ dialogue.chat_mode,
+ **{"chat_param": chat_param},
)
return chat
@@ -389,7 +419,7 @@ async def chat_prepare(dialogue: ConversationVo = Body(), user_token: UserReques
# dialogue.model_name = CFG.LLM_MODEL
logger.info(f"chat_prepare:{dialogue}")
## check conv_uid
- chat: BaseChat = get_chat_instance(dialogue, user_token.user_id)
+ chat: BaseChat = await get_chat_instance(dialogue, user_token.user_id)
if len(chat.history_message) > 0:
return Result.succ(None)
resp = await chat.prepare()
@@ -402,6 +432,15 @@ async def stream_error(msg):
@router.post("/v1/chat/completions")
async def chat_completions(dialogue: ConversationVo = Body(), user_token: UserRequest = Depends(get_user_from_headers)):
+ print(
+ f"chat_completions:{dialogue.chat_mode},{dialogue.select_param},{dialogue.model_name}"
+ )
+ with root_tracer.start_span(
+ "get_chat_instance", span_type=SpanType.CHAT, metadata=dialogue.dict()
+ ):
+ chat: BaseChat = await get_chat_instance(dialogue)
+ # background_tasks = BackgroundTasks()
+ # background_tasks.add_task(release_model_semaphore)
headers = {
"Content-Type": "text/event-stream",
"Cache-Control": "no-cache",
@@ -454,7 +493,7 @@ async def model_types(controller: BaseModelController = Depends(get_model_contro
return Result.succ(list(types))
except Exception as e:
- return Result.faild(code="E000X", msg=f"controller model types error {e}")
+ return Result.failed(code="E000X", msg=f"controller model types error {e}")
@router.get("/v1/model/supports")
@@ -464,7 +503,7 @@ async def model_supports(worker_manager: WorkerManager = Depends(get_worker_mana
models = await worker_manager.supported_models()
return Result.succ(FlatSupportedModel.from_supports(models))
except Exception as e:
- return Result.faild(code="E000X", msg=f"Fetch supportd models error {e}")
+ return Result.failed(code="E000X", msg=f"Fetch supportd models error {e}")
@router.get("/v1/github/callback")
@@ -498,8 +537,9 @@ async def github_access_token(code: str = None):
async def no_stream_generator(chat):
- msg = await chat.nostream_call()
- yield f"data: {msg}\n\n"
+ with root_tracer.start_span("no_stream_generator"):
+ msg = await chat.nostream_call()
+ yield f"data: {msg}\n\n"
async def stream_generator(chat, incremental: bool, model_name: str):
@@ -516,6 +556,7 @@ async def stream_generator(chat, incremental: bool, model_name: str):
Yields:
_type_: streaming responses
"""
+ span = root_tracer.start_span("stream_generator")
msg = "[LLM_ERROR]: llm server has no output, maybe your prompt template is wrong."
stream_id = f"chatcmpl-{str(uuid.uuid1())}"
@@ -541,6 +582,7 @@ async def stream_generator(chat, incremental: bool, model_name: str):
await asyncio.sleep(0.02)
if incremental:
yield "data: [DONE]\n\n"
+ span.end()
def message2Vo(message: dict, order, model_name) -> MessageVo:
diff --git a/pilot/openapi/api_v1/editor/api_editor_v1.py b/pilot/openapi/api_v1/editor/api_editor_v1.py
index e41998942..9653a1bad 100644
--- a/pilot/openapi/api_v1/editor/api_editor_v1.py
+++ b/pilot/openapi/api_v1/editor/api_editor_v1.py
@@ -107,7 +107,7 @@ async def get_editor_sql(con_uid: str, round: int):
.replace("\n", " ")
)
return Result.succ(json.loads(context))
- return Result.faild(msg="not have sql!")
+ return Result.failed(msg="not have sql!")
@router.post("/v1/editor/sql/run", response_model=Result[SqlRunData])
@@ -116,7 +116,7 @@ async def editor_sql_run(run_param: dict = Body()):
db_name = run_param["db_name"]
sql = run_param["sql"]
if not db_name and not sql:
- return Result.faild("SQL run param error!")
+ return Result.failed("SQL run param error!")
conn = CFG.LOCAL_DB_MANAGE.get_connect(db_name)
try:
@@ -134,6 +134,7 @@ async def editor_sql_run(run_param: dict = Body()):
)
return Result.succ(sql_run_data)
except Exception as e:
+ logging.error("editor_sql_run exception!" + str(e))
return Result.succ(
SqlRunData(result_info=str(e), run_cost=0, colunms=[], values=[])
)
@@ -165,11 +166,12 @@ async def sql_editor_submit(sql_edit_context: ChatSqlEditContext = Body()):
if element["type"] == "view":
data_loader = DbDataLoader()
element["data"]["content"] = data_loader.get_table_view_by_conn(
- conn.run(sql_edit_context.new_sql), sql_edit_context.new_speak
+ conn.run_to_df(sql_edit_context.new_sql),
+ sql_edit_context.new_speak,
)
history_mem.update(history_messages)
return Result.succ(None)
- return Result.faild(msg="Edit Faild!")
+ return Result.failed(msg="Edit Failed!")
@router.get("/v1/editor/chart/list", response_model=Result[ChartList])
@@ -191,7 +193,7 @@ async def get_editor_chart_list(con_uid: str):
charts=json.loads(element["data"]["content"]),
)
return Result.succ(chart_list)
- return Result.faild(msg="Not have charts!")
+ return Result.failed(msg="Not have charts!")
@router.post("/v1/editor/chart/info", response_model=Result[ChartDetail])
@@ -210,7 +212,7 @@ async def get_editor_chart_info(param: dict = Body()):
logger.error(
"this dashboard dialogue version too old, can't support editor!"
)
- return Result.faild(
+ return Result.failed(
msg="this dashboard dialogue version too old, can't support editor!"
)
for element in last_round["messages"]:
@@ -234,7 +236,7 @@ async def get_editor_chart_info(param: dict = Body()):
)
return Result.succ(detail)
- return Result.faild(msg="Can't Find Chart Detail Info!")
+ return Result.failed(msg="Can't Find Chart Detail Info!")
@router.post("/v1/editor/chart/run", response_model=Result[ChartRunData])
@@ -244,7 +246,7 @@ async def editor_chart_run(run_param: dict = Body()):
sql = run_param["sql"]
chart_type = run_param["chart_type"]
if not db_name and not sql:
- return Result.faild("SQL run param error!")
+ return Result.failed("SQL run param error!")
try:
dashboard_data_loader: DashboardDataLoader = DashboardDataLoader()
db_conn = CFG.LOCAL_DB_MANAGE.get_connect(db_name)
@@ -334,7 +336,7 @@ async def chart_editor_submit(chart_edit_context: ChatChartEditContext = Body())
)
except Exception as e:
logger.error(f"edit chart exception!{str(e)}", e)
- return Result.faild(msg=f"Edit chart exception!{str(e)}")
+ return Result.failed(msg=f"Edit chart exception!{str(e)}")
history_mem.update(history_messages)
return Result.succ(None)
- return Result.faild(msg="Edit Faild!")
+ return Result.failed(msg="Edit Failed!")
diff --git a/pilot/openapi/api_v1/feedback/feed_back_db.py b/pilot/openapi/api_v1/feedback/feed_back_db.py
index 3d697263b..02afb3215 100644
--- a/pilot/openapi/api_v1/feedback/feed_back_db.py
+++ b/pilot/openapi/api_v1/feedback/feed_back_db.py
@@ -3,7 +3,12 @@ from datetime import datetime
from sqlalchemy import Column, Integer, Text, String, DateTime
from pilot.base_modules.meta_data.base_dao import BaseDao
-from pilot.base_modules.meta_data.meta_data import Base, engine, session
+from pilot.base_modules.meta_data.meta_data import (
+ Base,
+ engine,
+ session,
+ META_DATA_DATABASE,
+)
from pilot.openapi.api_v1.feedback.feed_back_model import FeedBackBody
@@ -36,7 +41,10 @@ class ChatFeedBackEntity(Base):
class ChatFeedBackDao(BaseDao):
def __init__(self):
super().__init__(
- database="dbgpt", orm_base=Base, db_engine=engine, session=session
+ database=META_DATA_DATABASE,
+ orm_base=Base,
+ db_engine=engine,
+ session=session,
)
def create_or_update_chat_feed_back(self, feed_back: FeedBackBody):
diff --git a/pilot/openapi/api_view_model.py b/pilot/openapi/api_view_model.py
index 60065f2f2..af1aa4b9c 100644
--- a/pilot/openapi/api_view_model.py
+++ b/pilot/openapi/api_view_model.py
@@ -17,11 +17,11 @@ class Result(Generic[T], BaseModel):
return Result(success=True, err_code=None, err_msg=None, data=data)
@classmethod
- def faild(cls, msg):
+ def failed(cls, msg):
return Result(success=False, err_code="E000X", err_msg=msg, data=None)
@classmethod
- def faild(cls, code, msg):
+ def failed(cls, code, msg):
return Result(success=False, err_code=code, err_msg=msg, data=None)
diff --git a/pilot/openapi/base.py b/pilot/openapi/base.py
index 506254ec7..d8c814787 100644
--- a/pilot/openapi/base.py
+++ b/pilot/openapi/base.py
@@ -7,4 +7,4 @@ async def validation_exception_handler(request: Request, exc: RequestValidationE
message = ""
for error in exc.errors():
message += ".".join(error.get("loc")) + ":" + error.get("msg") + ";"
- return Result.faild(code="E0001", msg=message)
+ return Result.failed(code="E0001", msg=message)
diff --git a/pilot/out_parser/base.py b/pilot/out_parser/base.py
index 78286f3d3..a7268ab68 100644
--- a/pilot/out_parser/base.py
+++ b/pilot/out_parser/base.py
@@ -26,6 +26,10 @@ class BaseOutputParser(ABC):
def __init__(self, sep: str, is_stream_out: bool = True):
self.sep = sep
self.is_stream_out = is_stream_out
+ self.data_schema = None
+
+ def update(self, data_schema):
+ self.data_schema = data_schema
def __post_process_code(self, code):
sep = "\n```"
@@ -115,7 +119,9 @@ class BaseOutputParser(ABC):
print("un_stream ai response:", ai_response)
return ai_response
else:
- raise ValueError("Model server error!code=" + resp_obj_ex["error_code"])
+ raise ValueError(
+ f"""Model server error!code={resp_obj_ex["error_code"]}, errmsg is {resp_obj_ex["text"]}"""
+ )
def __illegal_json_ends(self, s):
temp_json = s
@@ -202,16 +208,23 @@ class BaseOutputParser(ABC):
if not cleaned_output.startswith("{") or not cleaned_output.endswith("}"):
logger.info("illegal json processing:\n" + cleaned_output)
cleaned_output = self.__extract_json(cleaned_output)
+
+ if not cleaned_output or len(cleaned_output) <= 0:
+ return model_out_text
+
cleaned_output = (
cleaned_output.strip()
.replace("\\n", " ")
.replace("\n", " ")
.replace("\\", " ")
+ .replace("\_", "_")
)
cleaned_output = self.__illegal_json_ends(cleaned_output)
return cleaned_output
- def parse_view_response(self, ai_text, data) -> str:
+ def parse_view_response(
+ self, ai_text, data, parse_prompt_response: Any = None
+ ) -> str:
"""
parse the ai response info to user view
Args:
@@ -242,7 +255,9 @@ class BaseOutputParser(ABC):
def _parse_model_response(response: ResponseTye):
- if isinstance(response, ModelOutput):
+ if response is None:
+ resp_obj_ex = ""
+ elif isinstance(response, ModelOutput):
resp_obj_ex = asdict(response)
elif isinstance(response, str):
resp_obj_ex = json.loads(response)
diff --git a/pilot/scene/base.py b/pilot/scene/base.py
index b56176991..e3478f7c3 100644
--- a/pilot/scene/base.py
+++ b/pilot/scene/base.py
@@ -82,6 +82,30 @@ class ChatScene(Enum):
"Dialogue through natural language and private documents and knowledge bases.",
["Knowledge Space Select"],
)
+ ExtractTriplet = Scene(
+ "extract_triplet",
+ "Extract Triplet",
+ "Extract Triplet",
+ ["Extract Select"],
+ True,
+ )
+ ExtractSummary = Scene(
+ "extract_summary",
+ "Extract Summary",
+ "Extract Summary",
+ ["Extract Select"],
+ True,
+ )
+ ExtractRefineSummary = Scene(
+ "extract_refine_summary",
+ "Extract Summary",
+ "Extract Summary",
+ ["Extract Select"],
+ True,
+ )
+ ExtractEntity = Scene(
+ "extract_entity", "Extract Entity", "Extract Entity", ["Extract Select"], True
+ )
@staticmethod
def of_mode(mode):
diff --git a/pilot/scene/base_chat.py b/pilot/scene/base_chat.py
index 571706b88..deb397dcf 100644
--- a/pilot/scene/base_chat.py
+++ b/pilot/scene/base_chat.py
@@ -12,8 +12,12 @@ from pilot.prompts.prompt_new import PromptTemplate
from pilot.scene.base_message import ModelMessage, ModelMessageRoleType
from pilot.scene.message import OnceConversation
from pilot.utils import get_or_create_event_loop
+from pilot.utils.executor_utils import ExecutorFactory, blocking_func_to_async
+from pilot.utils.tracer import root_tracer, trace
from pydantic import Extra
from pilot.memory.chat_history.chat_hisotry_factory import ChatHistory
+from pilot.awel import BaseOperator, SimpleCallDataInputSource, InputOperator, DAG
+from pilot.model.operator.model_operator import ModelOperator, ModelStreamOperator
logger = logging.getLogger(__name__)
headers = {"User-Agent": "dbgpt Client"}
@@ -37,6 +41,7 @@ class BaseChat(ABC):
arbitrary_types_allowed = True
+ @trace("BaseChat.__init__")
def __init__(self, chat_param: Dict):
"""Chat Module Initialization
Args:
@@ -53,6 +58,7 @@ class BaseChat(ABC):
chat_param["model_name"] if chat_param["model_name"] else CFG.LLM_MODEL
)
self.llm_echo = False
+ self.model_cache_enable = chat_param.get("model_cache_enable", False)
### load prompt template
# self.prompt_template: PromptTemplate = CFG.prompt_templates[
@@ -62,7 +68,7 @@ class BaseChat(ABC):
CFG.prompt_template_registry.get_prompt_template(
self.chat_mode.value(),
language=CFG.LANGUAGE,
- model_name=CFG.LLM_MODEL,
+ model_name=self.llm_model,
proxyllm_backend=CFG.PROXYLLM_BACKEND,
)
)
@@ -82,6 +88,15 @@ class BaseChat(ABC):
self.current_tokens_used: int = 0
if chat_param["user_id"]:
self.user_id = chat_param["user_id"]
+ # The executor to submit blocking function
+ self._executor = CFG.SYSTEM_APP.get_component(
+ ComponentType.EXECUTOR_DEFAULT, ExecutorFactory
+ ).create()
+
+ self._model_operator: BaseOperator = _build_model_operator()
+ self._model_stream_operator: BaseOperator = _build_model_operator(
+ is_stream=True, dag_name="llm_stream_model_dag"
+ )
class Config:
"""Configuration for this pydantic object."""
@@ -94,12 +109,21 @@ class BaseChat(ABC):
raise NotImplementedError("Not supported for this chat type.")
@abstractmethod
- def generate_input_values(self):
- pass
+ async def generate_input_values(self) -> Dict:
+ """Generate input to LLM
+
+ Please note that you must not perform any blocking operations in this function
+
+ Returns:
+ a dictionary to be formatted by prompt template
+ """
def do_action(self, prompt_response):
return prompt_response
+ def message_adjust(self):
+ pass
+
def get_llm_speak(self, prompt_define_response):
if hasattr(prompt_define_response, "thoughts"):
if isinstance(prompt_define_response.thoughts, dict):
@@ -118,18 +142,24 @@ class BaseChat(ABC):
speak_to_user = prompt_define_response
return speak_to_user
- def __call_base(self):
- input_values = self.generate_input_values()
+ async def __call_base(self):
+ input_values = await self.generate_input_values()
### Chat sequence advance
self.current_message.chat_order = len(self.history_message) + 1
self.current_message.add_user_message(self.current_user_input)
self.current_message.start_date = datetime.datetime.now().strftime(
"%Y-%m-%d %H:%M:%S"
)
-
self.current_message.tokens = 0
if self.prompt_template.template:
- current_prompt = self.prompt_template.format(**input_values)
+ metadata = {
+ "template_scene": self.prompt_template.template_scene,
+ "input_values": input_values,
+ }
+ with root_tracer.start_span(
+ "BaseChat.__call_base.prompt_template.format", metadata=metadata
+ ):
+ current_prompt = self.prompt_template.format(**input_values)
self.current_message.add_system_message(current_prompt)
llm_messages = self.generate_llm_messages()
@@ -137,14 +167,13 @@ class BaseChat(ABC):
# Not new server mode, we convert the message format(List[ModelMessage]) to list of dict
# fix the error of "Object of type ModelMessage is not JSON serializable" when passing the payload to request.post
llm_messages = list(map(lambda m: m.dict(), llm_messages))
-
payload = {
"model": self.llm_model,
"prompt": self.generate_llm_text(),
"messages": llm_messages,
"temperature": float(self.prompt_template.temperature),
"max_new_tokens": int(self.prompt_template.max_new_tokens),
- "stop": self.prompt_template.sep,
+ # "stop": self.prompt_template.sep,
"echo": self.llm_echo,
}
return payload
@@ -152,6 +181,9 @@ class BaseChat(ABC):
def stream_plugin_call(self, text):
return text
+ def stream_call_reinforce_fn(self, text):
+ return text
+
async def check_iterator_end(iterator):
try:
await asyncio.anext(iterator)
@@ -159,20 +191,30 @@ class BaseChat(ABC):
except StopAsyncIteration:
return True # 迭代器已经执行结束
+ def _get_span_metadata(self, payload: Dict) -> Dict:
+ metadata = {k: v for k, v in payload.items()}
+ del metadata["prompt"]
+ metadata["messages"] = list(
+ map(lambda m: m if isinstance(m, dict) else m.dict(), metadata["messages"])
+ )
+ return metadata
+
async def stream_call(self):
# TODO Retry when server connection error
- payload = self.__call_base()
+ payload = await self.__call_base()
self.skip_echo_len = len(payload.get("prompt").replace("", " ")) + 11
logger.info(f"Requert: \n{payload}")
ai_response_text = ""
+ span = root_tracer.start_span(
+ "BaseChat.stream_call", metadata=self._get_span_metadata(payload)
+ )
+ payload["span_id"] = span.span_id
+ payload["model_cache_enable"] = self.model_cache_enable
try:
- from pilot.model.cluster import WorkerManagerFactory
-
- worker_manager = CFG.SYSTEM_APP.get_component(
- ComponentType.WORKER_MANAGER_FACTORY, WorkerManagerFactory
- ).create()
- async for output in worker_manager.generate_stream(payload):
+ async for output in await self._model_stream_operator.call_stream(
+ call_data={"data": payload}
+ ):
### Plug-in research in result generation
msg = self.prompt_template.output_parser.parse_model_stream_resp_ex(
output, self.skip_echo_len
@@ -181,28 +223,33 @@ class BaseChat(ABC):
view_msg = view_msg.replace("\n", "\\n")
yield view_msg
self.current_message.add_ai_message(msg)
+ view_msg = self.stream_call_reinforce_fn(view_msg)
self.current_message.add_view_message(view_msg)
+ span.end()
except Exception as e:
print(traceback.format_exc())
- logger.error("model response parase faild!" + str(e))
+ logger.error("model response parase failed!" + str(e))
self.current_message.add_view_message(
f"""ERROR!{str(e)}\n {ai_response_text} """
)
### store current conversation
+ span.end(metadata={"error": str(e)})
self.memory.append(self.current_message, self.user_id)
async def nostream_call(self):
- payload = self.__call_base()
+ payload = await self.__call_base()
logger.info(f"Request: \n{payload}")
ai_response_text = ""
+ span = root_tracer.start_span(
+ "BaseChat.nostream_call", metadata=self._get_span_metadata(payload)
+ )
+ payload["span_id"] = span.span_id
+ payload["model_cache_enable"] = self.model_cache_enable
try:
- from pilot.model.cluster import WorkerManagerFactory
-
- worker_manager = CFG.SYSTEM_APP.get_component(
- ComponentType.WORKER_MANAGER_FACTORY, WorkerManagerFactory
- ).create()
-
- model_output = await worker_manager.generate(payload)
+ with root_tracer.start_span("BaseChat.invoke_worker_manager.generate"):
+ model_output = await self._model_operator.call(
+ call_data={"data": payload}
+ )
### output parse
ai_response_text = (
@@ -217,27 +264,78 @@ class BaseChat(ABC):
ai_response_text
)
)
- ### run
- result = self.do_action(prompt_define_response)
+ metadata = {
+ "model_output": model_output.to_dict(),
+ "ai_response_text": ai_response_text,
+ "prompt_define_response": self._parse_prompt_define_response(
+ prompt_define_response
+ ),
+ }
+ with root_tracer.start_span("BaseChat.do_action", metadata=metadata):
+ ### run
+ result = await blocking_func_to_async(
+ self._executor, self.do_action, prompt_define_response
+ )
### llm speaker
speak_to_user = self.get_llm_speak(prompt_define_response)
- view_message = self.prompt_template.output_parser.parse_view_response(
- speak_to_user, result
+ # view_message = self.prompt_template.output_parser.parse_view_response(
+ # speak_to_user, result
+ # )
+ view_message = await blocking_func_to_async(
+ self._executor,
+ self.prompt_template.output_parser.parse_view_response,
+ speak_to_user,
+ result,
+ prompt_define_response,
)
+
view_message = view_message.replace("\n", "\\n")
self.current_message.add_view_message(view_message)
+ self.message_adjust()
+
+ span.end()
except Exception as e:
print(traceback.format_exc())
logger.error("model response parase faild!" + str(e))
self.current_message.add_view_message(
f"""ERROR!{str(e)}\n {ai_response_text} """
)
+ span.end(metadata={"error": str(e)})
### store dialogue
self.memory.append(self.current_message, self.user_id)
return self.current_ai_response()
+ async def get_llm_response(self):
+ payload = await self.__call_base()
+ logger.info(f"Request: \n{payload}")
+ ai_response_text = ""
+ payload["model_cache_enable"] = self.model_cache_enable
+ try:
+ model_output = await self._model_operator.call(call_data={"data": payload})
+ ### output parse
+ ai_response_text = (
+ self.prompt_template.output_parser.parse_model_nostream_resp(
+ model_output, self.prompt_template.sep
+ )
+ )
+ ### model result deal
+ self.current_message.add_ai_message(ai_response_text)
+ prompt_define_response = None
+ prompt_define_response = (
+ self.prompt_template.output_parser.parse_prompt_response(
+ ai_response_text
+ )
+ )
+ except Exception as e:
+ print(traceback.format_exc())
+ logger.error("model response parse failed!" + str(e))
+ self.current_message.add_view_message(
+ f"""model response parse failed!{str(e)}\n {ai_response_text} """
+ )
+ return prompt_define_response
+
def _blocking_stream_call(self):
logger.warn(
"_blocking_stream_call is only temporarily used in webserver and will be deleted soon, please use stream_call to replace it for higher performance"
@@ -277,16 +375,18 @@ class BaseChat(ABC):
if self.prompt_template.template_define:
text += self.prompt_template.template_define + self.prompt_template.sep
### Load prompt
- text += self.__load_system_message()
+ text += _load_system_message(self.current_message, self.prompt_template)
### Load examples
- text += self.__load_example_messages()
+ text += _load_example_messages(self.prompt_template)
### Load History
- text += self.__load_histroy_messages()
+ text += _load_history_messages(
+ self.prompt_template, self.history_message, self.chat_retention_rounds
+ )
### Load User Input
- text += self.__load_user_message()
+ text += _load_user_message(self.current_message, self.prompt_template)
return text
def generate_llm_messages(self) -> List[ModelMessage]:
@@ -304,137 +404,26 @@ class BaseChat(ABC):
)
)
### Load prompt
- messages += self.__load_system_message(str_message=False)
+ messages += _load_system_message(
+ self.current_message, self.prompt_template, str_message=False
+ )
### Load examples
- messages += self.__load_example_messages(str_message=False)
+ messages += _load_example_messages(self.prompt_template, str_message=False)
### Load History
- messages += self.__load_histroy_messages(str_message=False)
+ messages += _load_history_messages(
+ self.prompt_template,
+ self.history_message,
+ self.chat_retention_rounds,
+ str_message=False,
+ )
### Load User Input
- messages += self.__load_user_message(str_message=False)
+ messages += _load_user_message(
+ self.current_message, self.prompt_template, str_message=False
+ )
return messages
- def __load_system_message(self, str_message: bool = True):
- system_convs = self.current_message.get_system_conv()
- system_text = ""
- system_messages = []
- for system_conv in system_convs:
- system_text += (
- system_conv.type + ":" + system_conv.content + self.prompt_template.sep
- )
- system_messages.append(
- ModelMessage(role=system_conv.type, content=system_conv.content)
- )
- return system_text if str_message else system_messages
-
- def __load_user_message(self, str_message: bool = True):
- user_conv = self.current_message.get_user_conv()
- user_messages = []
- if user_conv:
- user_text = (
- user_conv.type + ":" + user_conv.content + self.prompt_template.sep
- )
- user_messages.append(
- ModelMessage(role=user_conv.type, content=user_conv.content)
- )
- return user_text if str_message else user_messages
- else:
- raise ValueError("Hi! What do you want to talk about?")
-
- def __load_example_messages(self, str_message: bool = True):
- example_text = ""
- example_messages = []
- if self.prompt_template.example_selector:
- for round_conv in self.prompt_template.example_selector.examples():
- for round_message in round_conv["messages"]:
- if not round_message["type"] in [
- ModelMessageRoleType.VIEW,
- ModelMessageRoleType.SYSTEM,
- ]:
- message_type = round_message["type"]
- message_content = round_message["data"]["content"]
- example_text += (
- message_type
- + ":"
- + message_content
- + self.prompt_template.sep
- )
- example_messages.append(
- ModelMessage(role=message_type, content=message_content)
- )
- return example_text if str_message else example_messages
-
- def __load_histroy_messages(self, str_message: bool = True):
- history_text = ""
- history_messages = []
- if self.prompt_template.need_historical_messages:
- if self.history_message:
- logger.info(
- f"There are already {len(self.history_message)} rounds of conversations! Will use {self.chat_retention_rounds} rounds of content as history!"
- )
- if len(self.history_message) > self.chat_retention_rounds:
- for first_message in self.history_message[0]["messages"]:
- if not first_message["type"] in [
- ModelMessageRoleType.VIEW,
- ModelMessageRoleType.SYSTEM,
- ]:
- message_type = first_message["type"]
- message_content = first_message["data"]["content"]
- history_text += (
- message_type
- + ":"
- + message_content
- + self.prompt_template.sep
- )
- history_messages.append(
- ModelMessage(role=message_type, content=message_content)
- )
- if self.chat_retention_rounds > 1:
- index = self.chat_retention_rounds - 1
- for round_conv in self.history_message[-index:]:
- for round_message in round_conv["messages"]:
- if not round_message["type"] in [
- ModelMessageRoleType.VIEW,
- ModelMessageRoleType.SYSTEM,
- ]:
- message_type = round_message["type"]
- message_content = round_message["data"]["content"]
- history_text += (
- message_type
- + ":"
- + message_content
- + self.prompt_template.sep
- )
- history_messages.append(
- ModelMessage(
- role=message_type, content=message_content
- )
- )
-
- else:
- ### user all history
- for conversation in self.history_message:
- for message in conversation["messages"]:
- ### histroy message not have promot and view info
- if not message["type"] in [
- ModelMessageRoleType.VIEW,
- ModelMessageRoleType.SYSTEM,
- ]:
- message_type = message["type"]
- message_content = message["data"]["content"]
- history_text += (
- message_type
- + ":"
- + message_content
- + self.prompt_template.sep
- )
- history_messages.append(
- ModelMessage(role=message_type, content=message_content)
- )
-
- return history_text if str_message else history_messages
-
def current_ai_response(self) -> str:
for message in self.current_message.messages:
if message.type == "view":
@@ -451,3 +440,230 @@ class BaseChat(ABC):
"""
pass
+
+ def _parse_prompt_define_response(self, prompt_define_response: Any) -> Any:
+ if not prompt_define_response:
+ return ""
+ if isinstance(prompt_define_response, str) or isinstance(
+ prompt_define_response, dict
+ ):
+ return prompt_define_response
+ if isinstance(prompt_define_response, tuple):
+ if hasattr(prompt_define_response, "_asdict"):
+ # namedtuple
+ return prompt_define_response._asdict()
+ else:
+ return dict(
+ zip(range(len(prompt_define_response)), prompt_define_response)
+ )
+ else:
+ return prompt_define_response
+
+
+def _build_model_operator(
+ is_stream: bool = False, dag_name: str = "llm_model_dag"
+) -> BaseOperator:
+ """Builds and returns a model processing workflow (DAG) operator.
+
+ This function constructs a Directed Acyclic Graph (DAG) for processing data using a model.
+ It includes caching and branching logic to either fetch results from a cache or process
+ data using the model. It supports both streaming and non-streaming modes.
+
+ .. code-block:: python
+ input_node >> cache_check_branch_node
+ cache_check_branch_node >> model_node >> save_cached_node >> join_node
+ cache_check_branch_node >> cached_node >> join_node
+
+ equivalent to::
+
+ -> model_node -> save_cached_node ->
+ / \
+ input_node -> cache_check_branch_node ---> join_node
+ \ /
+ -> cached_node ------------------- ->
+
+ Args:
+ is_stream (bool): Flag to determine if the operator should process data in streaming mode.
+ dag_name (str): Name of the DAG.
+
+ Returns:
+ BaseOperator: The final operator in the constructed DAG, typically a join node.
+ """
+ from pilot.model.cluster import WorkerManagerFactory
+ from pilot.awel import JoinOperator
+ from pilot.model.operator.model_operator import (
+ ModelCacheBranchOperator,
+ CachedModelStreamOperator,
+ CachedModelOperator,
+ ModelSaveCacheOperator,
+ ModelStreamSaveCacheOperator,
+ )
+ from pilot.cache import CacheManager
+
+ # Fetch worker and cache managers from the system configuration
+ worker_manager = CFG.SYSTEM_APP.get_component(
+ ComponentType.WORKER_MANAGER_FACTORY, WorkerManagerFactory
+ ).create()
+ cache_manager: CacheManager = CFG.SYSTEM_APP.get_component(
+ ComponentType.MODEL_CACHE_MANAGER, CacheManager
+ )
+ # Define task names for the model and cache nodes
+ model_task_name = "llm_model_node"
+ cache_task_name = "llm_model_cache_node"
+
+ with DAG(dag_name):
+ # Create an input node
+ input_node = InputOperator(SimpleCallDataInputSource())
+ # Determine if the workflow should operate in streaming mode
+ if is_stream:
+ model_node = ModelStreamOperator(worker_manager, task_name=model_task_name)
+ cached_node = CachedModelStreamOperator(
+ cache_manager, task_name=cache_task_name
+ )
+ save_cached_node = ModelStreamSaveCacheOperator(cache_manager)
+ else:
+ model_node = ModelOperator(worker_manager, task_name=model_task_name)
+ cached_node = CachedModelOperator(cache_manager, task_name=cache_task_name)
+ save_cached_node = ModelSaveCacheOperator(cache_manager)
+
+ # Create a branch node to decide between fetching from cache or processing with the model
+ cache_check_branch_node = ModelCacheBranchOperator(
+ cache_manager,
+ model_task_name="llm_model_node",
+ cache_task_name="llm_model_cache_node",
+ )
+ # Create a join node to merge outputs from the model and cache nodes, just keep the first not empty output
+ join_node = JoinOperator(
+ combine_function=lambda model_out, cache_out: cache_out or model_out
+ )
+
+ # Define the workflow structure using the >> operator
+ input_node >> cache_check_branch_node
+ cache_check_branch_node >> model_node >> save_cached_node >> join_node
+ cache_check_branch_node >> cached_node >> join_node
+
+ return join_node
+
+
+def _load_system_message(
+ current_message: OnceConversation,
+ prompt_template: PromptTemplate,
+ str_message: bool = True,
+):
+ system_convs = current_message.get_system_conv()
+ system_text = ""
+ system_messages = []
+ for system_conv in system_convs:
+ system_text += (
+ system_conv.type + ":" + system_conv.content + prompt_template.sep
+ )
+ system_messages.append(
+ ModelMessage(role=system_conv.type, content=system_conv.content)
+ )
+ return system_text if str_message else system_messages
+
+
+def _load_user_message(
+ current_message: OnceConversation,
+ prompt_template: PromptTemplate,
+ str_message: bool = True,
+):
+ user_conv = current_message.get_user_conv()
+ user_messages = []
+ if user_conv:
+ user_text = user_conv.type + ":" + user_conv.content + prompt_template.sep
+ user_messages.append(
+ ModelMessage(role=user_conv.type, content=user_conv.content)
+ )
+ return user_text if str_message else user_messages
+ else:
+ raise ValueError("Hi! What do you want to talk about?")
+
+
+def _load_example_messages(prompt_template: PromptTemplate, str_message: bool = True):
+ example_text = ""
+ example_messages = []
+ if prompt_template.example_selector:
+ for round_conv in prompt_template.example_selector.examples():
+ for round_message in round_conv["messages"]:
+ if not round_message["type"] in [
+ ModelMessageRoleType.VIEW,
+ ModelMessageRoleType.SYSTEM,
+ ]:
+ message_type = round_message["type"]
+ message_content = round_message["data"]["content"]
+ example_text += (
+ message_type + ":" + message_content + prompt_template.sep
+ )
+ example_messages.append(
+ ModelMessage(role=message_type, content=message_content)
+ )
+ return example_text if str_message else example_messages
+
+
+def _load_history_messages(
+ prompt_template: PromptTemplate,
+ history_message: List[OnceConversation],
+ chat_retention_rounds: int,
+ str_message: bool = True,
+):
+ history_text = ""
+ history_messages = []
+ if prompt_template.need_historical_messages:
+ if history_message:
+ logger.info(
+ f"There are already {len(history_message)} rounds of conversations! Will use {chat_retention_rounds} rounds of content as history!"
+ )
+ if len(history_message) > chat_retention_rounds:
+ for first_message in history_message[0]["messages"]:
+ if not first_message["type"] in [
+ ModelMessageRoleType.VIEW,
+ ModelMessageRoleType.SYSTEM,
+ ]:
+ message_type = first_message["type"]
+ message_content = first_message["data"]["content"]
+ history_text += (
+ message_type + ":" + message_content + prompt_template.sep
+ )
+ history_messages.append(
+ ModelMessage(role=message_type, content=message_content)
+ )
+ if chat_retention_rounds > 1:
+ index = chat_retention_rounds - 1
+ for round_conv in history_message[-index:]:
+ for round_message in round_conv["messages"]:
+ if not round_message["type"] in [
+ ModelMessageRoleType.VIEW,
+ ModelMessageRoleType.SYSTEM,
+ ]:
+ message_type = round_message["type"]
+ message_content = round_message["data"]["content"]
+ history_text += (
+ message_type
+ + ":"
+ + message_content
+ + prompt_template.sep
+ )
+ history_messages.append(
+ ModelMessage(role=message_type, content=message_content)
+ )
+
+ else:
+ ### user all history
+ for conversation in history_message:
+ for message in conversation["messages"]:
+ ### histroy message not have promot and view info
+ if not message["type"] in [
+ ModelMessageRoleType.VIEW,
+ ModelMessageRoleType.SYSTEM,
+ ]:
+ message_type = message["type"]
+ message_content = message["data"]["content"]
+ history_text += (
+ message_type + ":" + message_content + prompt_template.sep
+ )
+ history_messages.append(
+ ModelMessage(role=message_type, content=message_content)
+ )
+
+ return history_text if str_message else history_messages
diff --git a/pilot/scene/base_message.py b/pilot/scene/base_message.py
index eeb42a285..c4c10459c 100644
--- a/pilot/scene/base_message.py
+++ b/pilot/scene/base_message.py
@@ -1,7 +1,7 @@
from __future__ import annotations
from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Tuple, Optional
+from typing import Any, Dict, List, Tuple, Optional, Union
from pydantic import BaseModel, Field, root_validator
@@ -70,14 +70,6 @@ class SystemMessage(BaseMessage):
return "system"
-class ModelMessage(BaseModel):
- """Type of message that interaction between dbgpt-server and llm-server"""
-
- """Similar to openai's message format"""
- role: str
- content: str
-
-
class ModelMessageRoleType:
""" "Type of ModelMessage role"""
@@ -87,6 +79,77 @@ class ModelMessageRoleType:
VIEW = "view"
+class ModelMessage(BaseModel):
+ """Type of message that interaction between dbgpt-server and llm-server"""
+
+ """Similar to openai's message format"""
+ role: str
+ content: str
+
+ @staticmethod
+ def from_openai_messages(
+ messages: Union[str, List[Dict[str, str]]]
+ ) -> List["ModelMessage"]:
+ """Openai message format to current ModelMessage format"""
+ if isinstance(messages, str):
+ return [ModelMessage(role=ModelMessageRoleType.HUMAN, content=messages)]
+ result = []
+ for message in messages:
+ msg_role = message["role"]
+ content = message["content"]
+ if msg_role == "system":
+ result.append(
+ ModelMessage(role=ModelMessageRoleType.SYSTEM, content=content)
+ )
+ elif msg_role == "user":
+ result.append(
+ ModelMessage(role=ModelMessageRoleType.HUMAN, content=content)
+ )
+ elif msg_role == "assistant":
+ result.append(
+ ModelMessage(role=ModelMessageRoleType.AI, content=content)
+ )
+ else:
+ raise ValueError(f"Unknown role: {msg_role}")
+ return result
+
+ @staticmethod
+ def to_openai_messages(messages: List["ModelMessage"]) -> List[Dict[str, str]]:
+ """Convert to OpenAI message format and
+ hugggingface [Templates of Chat Models](https://huggingface.co/docs/transformers/v4.34.1/en/chat_templating)
+ """
+ history = []
+ # Add history conversation
+ for message in messages:
+ if message.role == ModelMessageRoleType.HUMAN:
+ history.append({"role": "user", "content": message.content})
+ elif message.role == ModelMessageRoleType.SYSTEM:
+ history.append({"role": "system", "content": message.content})
+ elif message.role == ModelMessageRoleType.AI:
+ history.append({"role": "assistant", "content": message.content})
+ else:
+ pass
+ # Move the last user's information to the end
+ temp_his = history[::-1]
+ last_user_input = None
+ for m in temp_his:
+ if m["role"] == "user":
+ last_user_input = m
+ break
+ if last_user_input:
+ history.remove(last_user_input)
+ history.append(last_user_input)
+ return history
+
+ @staticmethod
+ def to_dict_list(messages: List["ModelMessage"]) -> List[Dict[str, str]]:
+ return list(map(lambda m: m.dict(), messages))
+
+ @staticmethod
+ def build_human_message(content: str) -> "ModelMessage":
+ return ModelMessage(role=ModelMessageRoleType.HUMAN, content=content)
+
+
class Generation(BaseModel):
"""Output of a single generation."""
diff --git a/pilot/scene/chat_agent/chat.py b/pilot/scene/chat_agent/chat.py
index 0fc5a0375..09968d8f3 100644
--- a/pilot/scene/chat_agent/chat.py
+++ b/pilot/scene/chat_agent/chat.py
@@ -11,6 +11,7 @@ from pilot.common.string_utils import extract_content
from .prompt import prompt
from pilot.component import ComponentType
from pilot.base_modules.agent.controller import ModuleAgent
+from pilot.utils.tracer import root_tracer, trace
CFG = Config()
@@ -51,7 +52,8 @@ class ChatAgent(BaseChat):
self.api_call = ApiCall(plugin_generator=self.plugins_prompt_generator)
- def generate_input_values(self):
+ @trace()
+ async def generate_input_values(self) -> Dict[str, str]:
input_values = {
"user_goal": self.current_user_input,
"expand_constraints": self.__list_to_prompt_str(
@@ -62,8 +64,16 @@ class ChatAgent(BaseChat):
return input_values
def stream_plugin_call(self, text):
- text = text.replace("\n", " ")
- return self.api_call.run(text)
+ text = (
+ text.replace("\\n", " ")
+ .replace("\n", " ")
+ .replace("\_", "_")
+ .replace("\\", " ")
+ )
+ with root_tracer.start_span(
+ "ChatAgent.stream_plugin_call.api_call", metadata={"text": text}
+ ):
+ return self.api_call.run(text)
def __list_to_prompt_str(self, list: List) -> str:
return "\n".join(f"{i + 1 + 1}. {item}" for i, item in enumerate(list))
diff --git a/pilot/scene/chat_agent/out_parser.py b/pilot/scene/chat_agent/out_parser.py
index ecea328bb..962a40fbc 100644
--- a/pilot/scene/chat_agent/out_parser.py
+++ b/pilot/scene/chat_agent/out_parser.py
@@ -1,11 +1,6 @@
import json
from typing import Dict, NamedTuple
-from pilot.utils import build_logger
from pilot.out_parser.base import BaseOutputParser, T
-from pilot.configs.model_config import LOGDIR
-
-
-logger = build_logger("webserver", LOGDIR + "DbChatOutputParser.log")
class PluginAction(NamedTuple):
@@ -15,7 +10,7 @@ class PluginAction(NamedTuple):
class PluginChatOutputParser(BaseOutputParser):
- def parse_view_response(self, speak, data) -> str:
+ def parse_view_response(self, speak, data, prompt_response) -> str:
### tool out data to table view
print(f"parse_view_response:{speak},{str(data)}")
view_text = f"##### {speak}" + "\n" + str(data)
diff --git a/pilot/scene/chat_agent/prompt.py b/pilot/scene/chat_agent/prompt.py
index 7d01bfd2f..94151544a 100644
--- a/pilot/scene/chat_agent/prompt.py
+++ b/pilot/scene/chat_agent/prompt.py
@@ -42,7 +42,8 @@ _DEFAULT_TEMPLATE_ZH = """
3.根据上面约束的方式生成每个工具的调用,对于工具使用的提示文本,需要在工具使用前生成
4.如果用户目标无法理解和意图不明确,优先使用搜索引擎工具
5.参数内容可能需要根据用户的目标推理得到,不仅仅是从文本提取
- 6.约束条件和工具信息作为推理过程的辅助信息,不要表达在给用户的输出内容中
+ 6.约束条件和工具信息作为推理过程的辅助信息,对应内容不要表达在给用户的输出内容中
+ 7.不要把[INST] <", "").replace("", ""))
@@ -162,8 +231,12 @@ class LLMModelAdaper:
model_context["has_format_prompt"] = True
params["prompt"] = new_prompt
- # Overwrite model params:
- params["stop"] = conv.stop_str
+ custom_stop = params.get("stop")
+ custom_stop_token_ids = params.get("stop_token_ids")
+
+ # Prefer the value passed in from the input parameter
+ params["stop"] = custom_stop or conv_stop_str
+ params["stop_token_ids"] = custom_stop_token_ids or conv_stop_token_ids
return params, model_context
@@ -216,9 +289,16 @@ class FastChatLLMModelAdaperWrapper(LLMModelAdaper):
return self._adapter.load_model(model_path, from_pretrained_kwargs)
def get_generate_stream_function(self, model: "TorchNNModule", model_path: str):
- from fastchat.model.model_adapter import get_generate_stream_function
+ if _IS_BENCHMARK:
+ from pilot.utils.benchmarks.llm.fastchat_benchmarks_inference import (
+ generate_stream,
+ )
- return get_generate_stream_function(model, model_path)
+ return generate_stream
+ else:
+ from fastchat.model.model_adapter import get_generate_stream_function
+
+ return get_generate_stream_function(model, model_path)
def get_default_conv_template(
self, model_name: str, model_path: str
@@ -233,6 +313,69 @@ class FastChatLLMModelAdaperWrapper(LLMModelAdaper):
)
+class NewHFChatModelAdapter(LLMModelAdaper):
+ def load(self, model_path: str, from_pretrained_kwargs: dict):
+ try:
+ import transformers
+ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
+ except ImportError as exc:
+ raise ValueError(
+ "Could not import depend python package "
+ "Please install it with `pip install transformers`."
+ ) from exc
+ if not transformers.__version__ >= "4.34.0":
+ raise ValueError(
+ "Current model (Load by HFNewChatAdapter) require transformers.__version__>=4.34.0"
+ )
+ revision = from_pretrained_kwargs.get("revision", "main")
+ try:
+ tokenizer = AutoTokenizer.from_pretrained(
+ model_path,
+ use_fast=self.use_fast_tokenizer,
+ revision=revision,
+ trust_remote_code=True,
+ )
+ except TypeError:
+ tokenizer = AutoTokenizer.from_pretrained(
+ model_path, use_fast=False, revision=revision, trust_remote_code=True
+ )
+ try:
+ model = AutoModelForCausalLM.from_pretrained(
+ model_path, low_cpu_mem_usage=True, **from_pretrained_kwargs
+ )
+ except NameError:
+ model = AutoModel.from_pretrained(
+ model_path, low_cpu_mem_usage=True, **from_pretrained_kwargs
+ )
+ # tokenizer.use_default_system_prompt = False
+ return model, tokenizer
+
+ def get_generate_stream_function(self, model, model_path: str):
+ """Get the generate stream function of the model"""
+ from pilot.model.llm_out.hf_chat_llm import huggingface_chat_generate_stream
+
+ return huggingface_chat_generate_stream
+
+ def get_str_prompt(
+ self,
+ params: Dict,
+ messages: List[ModelMessage],
+ tokenizer: Any,
+ prompt_template: str = None,
+ ) -> Optional[str]:
+ from transformers import AutoTokenizer
+
+ if not tokenizer:
+ raise ValueError("tokenizer is is None")
+ tokenizer: AutoTokenizer = tokenizer
+
+ messages = ModelMessage.to_openai_messages(messages)
+ str_prompt = tokenizer.apply_chat_template(
+ messages, tokenize=False, add_generation_prompt=True
+ )
+ return str_prompt
+
+
def get_conv_template(name: str) -> "Conversation":
"""Get a conversation template."""
from fastchat.conversation import get_conv_template
@@ -261,6 +404,11 @@ def get_llm_model_adapter(
logger.info("Current model type is vllm, return VLLMModelAdaperWrapper")
return VLLMModelAdaperWrapper()
+ use_new_hf_chat_models = any(m in model_name.lower() for m in _NEW_HF_CHAT_MODELS)
+ if use_new_hf_chat_models:
+ logger.info(f"Current model {model_name} use NewHFChatModelAdapter")
+ return NewHFChatModelAdapter()
+
must_use_old = any(m in model_name for m in _OLD_MODELS)
if use_fastchat and not must_use_old:
logger.info("Use fastcat adapter")
@@ -297,6 +445,7 @@ def _get_fastchat_model_adapter(
if use_fastchat_monkey_patch:
model_adapter.get_model_adapter = _fastchat_get_adapter_monkey_patch
thread_local.model_name = model_name
+ _remove_black_list_model_of_fastchat()
if caller:
return caller(model_path)
finally:
@@ -340,6 +489,24 @@ def _fastchat_get_adapter_monkey_patch(model_path: str, model_name: str = None):
)
+@cache
+def _remove_black_list_model_of_fastchat():
+ from fastchat.model.model_adapter import model_adapters
+
+ black_list_models = []
+ for adapter in model_adapters:
+ try:
+ if (
+ adapter.get_default_conv_template("/data/not_exist_model_path").name
+ in _BLACK_LIST_MODLE_PROMPT
+ ):
+ black_list_models.append(adapter)
+ except Exception:
+ pass
+ for adapter in black_list_models:
+ model_adapters.remove(adapter)
+
+
def _dynamic_model_parser() -> Callable[[None], List[Type]]:
from pilot.utils.parameter_utils import _SimpleArgParser
from pilot.model.parameter import (
@@ -444,17 +611,50 @@ class VLLMModelAdaperWrapper(LLMModelAdaper):
# Covering the configuration of fastcaht, we will regularly feedback the code here to fastchat.
# We also recommend that you modify it directly in the fastchat repository.
+
+# source: https://huggingface.co/BAAI/AquilaChat2-34B/blob/4608b75855334b93329a771aee03869dbf7d88cc/predict.py#L212
register_conv_template(
Conversation(
- name="internlm-chat",
- system_message="A chat between a curious <|User|> and an <|Bot|>. The <|Bot|> gives helpful, detailed, and polite answers to the <|User|>'s questions.\n\n",
- roles=("<|User|>", "<|Bot|>"),
- sep_style=SeparatorStyle.CHATINTERN,
- sep="",
+ stop_str=["404
This page could not be found.
404
This page could not be found.
404
This page could not be found.
404
This page could not be found.