diff --git a/.env.template b/.env.template index 2058c36a3..f53bf31f7 100644 --- a/.env.template +++ b/.env.template @@ -54,17 +54,16 @@ KNOWLEDGE_SEARCH_TOP_SIZE=5 #*******************************************************************# #** DATABASE SETTINGS **# #*******************************************************************# -### MYSQL database(Current default database) -LOCAL_DB_TYPE=mysql -LOCAL_DB_USER=root -LOCAL_DB_PASSWORD=aa12345678 -LOCAL_DB_HOST=127.0.0.1 -LOCAL_DB_PORT=3306 - -### SQLite database (TODO: SQLite database will become the default database configuration when it is stable.) -# LOCAL_DB_PATH=data/default_sqlite.db -# LOCAL_DB_TYPE=sqlite +### SQLite database (Current default database) +LOCAL_DB_PATH=data/default_sqlite.db +LOCAL_DB_TYPE=sqlite +### MYSQL database +# LOCAL_DB_TYPE=mysql +# LOCAL_DB_USER=root +# LOCAL_DB_PASSWORD=aa12345678 +# LOCAL_DB_HOST=127.0.0.1 +# LOCAL_DB_PORT=3306 ### MILVUS ## MILVUS_ADDR - Milvus remote address (e.g. localhost:19530) diff --git a/docker/allinone/Dockerfile b/docker/allinone/Dockerfile index 496f7b2bf..48c8362cf 100644 --- a/docker/allinone/Dockerfile +++ b/docker/allinone/Dockerfile @@ -1,4 +1,4 @@ -ARG BASE_IMAGE="db-gpt:latest" +ARG BASE_IMAGE="eosphorosai/dbgpt:latest" FROM ${BASE_IMAGE} @@ -25,6 +25,6 @@ ENV LOCAL_DB_PASSWORD="$MYSQL_ROOT_PASSWORD" RUN cp /app/assets/schema/knowledge_management.sql /docker-entrypoint-initdb.d/ COPY docker/allinone/allinone-entrypoint.sh /usr/local/bin/allinone-entrypoint.sh -COPY docker/examples/sqls/ /docker-entrypoint-initdb.d/ +COPY docker/examples/sqls/*_mysql.sql /docker-entrypoint-initdb.d/ ENTRYPOINT ["/usr/local/bin/allinone-entrypoint.sh"] \ No newline at end of file diff --git a/docker/allinone/build_image.sh b/docker/allinone/build_image.sh index 9ff255cc8..b72a3a38c 100755 --- a/docker/allinone/build_image.sh +++ b/docker/allinone/build_image.sh @@ -4,6 +4,6 @@ SCRIPT_LOCATION=$0 cd "$(dirname "$SCRIPT_LOCATION")" WORK_DIR=$(pwd) -IMAGE_NAME="db-gpt-allinone" +IMAGE_NAME="eosphorosai/dbgpt-allinone" docker build -f Dockerfile -t $IMAGE_NAME $WORK_DIR/../../ \ No newline at end of file diff --git a/docker/allinone/run.sh b/docker/allinone/run.sh index 913dcef4d..0b015ec0b 100755 --- a/docker/allinone/run.sh +++ b/docker/allinone/run.sh @@ -1,6 +1,6 @@ #!/bin/bash -docker run --gpus "device=0" -d -p 3306:3306 \ +docker run --gpus all -d -p 3306:3306 \ -p 5000:5000 \ -e LOCAL_DB_HOST=127.0.0.1 \ -e LOCAL_DB_PASSWORD=aa123456 \ @@ -9,5 +9,5 @@ docker run --gpus "device=0" -d -p 3306:3306 \ -e LANGUAGE=zh \ -v /data:/data \ -v /data/models:/app/models \ - --name db-gpt-allinone \ - db-gpt-allinone \ No newline at end of file + --name dbgpt-allinone \ + eosphorosai/dbgpt-allinone \ No newline at end of file diff --git a/docker/allinone/run_proxyllm.sh b/docker/allinone/run_proxyllm.sh index 7823c5d78..239b7ce7a 100755 --- a/docker/allinone/run_proxyllm.sh +++ b/docker/allinone/run_proxyllm.sh @@ -4,7 +4,7 @@ PROXY_API_KEY="$PROXY_API_KEY" PROXY_SERVER_URL="${PROXY_SERVER_URL-'https://api.openai.com/v1/chat/completions'}" -docker run --gpus "device=0" -d -p 3306:3306 \ +docker run --gpus all -d -p 3306:3306 \ -p 5000:5000 \ -e LOCAL_DB_HOST=127.0.0.1 \ -e LOCAL_DB_PASSWORD=aa123456 \ @@ -15,5 +15,5 @@ docker run --gpus "device=0" -d -p 3306:3306 \ -e LANGUAGE=zh \ -v /data:/data \ -v /data/models:/app/models \ - --name db-gpt-allinone \ - db-gpt-allinone \ No newline at end of file + --name dbgpt-allinone \ + eosphorosai/dbgpt-allinone \ No newline at end of file diff --git a/docker/base/Dockerfile b/docker/base/Dockerfile index d65fc2a98..7df9e2d4e 100644 --- a/docker/base/Dockerfile +++ b/docker/base/Dockerfile @@ -45,4 +45,14 @@ RUN (if [ "${BUILD_LOCAL_CODE}" = "true" ]; \ else rm -rf /tmp/app; \ fi;) -EXPOSE 5000 \ No newline at end of file +ARG LOAD_EXAMPLES="true" + +RUN (if [ "${LOAD_EXAMPLES}" = "true" ]; \ + then mkdir -p /app/pilot/data && sqlite3 /app/pilot/data/default_sqlite.db < /app/docker/examples/sqls/case_1_student_manager_sqlite.sql \ + && sqlite3 /app/pilot/data/default_sqlite.db < /app/docker/examples/sqls/case_2_ecom_sqlite.sql \ + && sqlite3 /app/pilot/data/default_sqlite.db < /app/docker/examples/sqls/test_case_info_sqlite.sql; \ + fi;) + +EXPOSE 5000 + +CMD ["python3", "pilot/server/dbgpt_server.py"] \ No newline at end of file diff --git a/docker/base/build_image.sh b/docker/base/build_image.sh index 83d6ca18d..101dfebad 100755 --- a/docker/base/build_image.sh +++ b/docker/base/build_image.sh @@ -5,12 +5,13 @@ cd "$(dirname "$SCRIPT_LOCATION")" WORK_DIR=$(pwd) BASE_IMAGE="nvidia/cuda:11.8.0-devel-ubuntu22.04" -IMAGE_NAME="db-gpt" +IMAGE_NAME="eosphorosai/dbgpt" # zh: https://pypi.tuna.tsinghua.edu.cn/simple PIP_INDEX_URL="https://pypi.org/simple" # en or zh LANGUAGE="en" BUILD_LOCAL_CODE="false" +LOAD_EXAMPLES="true" usage () { echo "USAGE: $0 [--base-image nvidia/cuda:11.8.0-devel-ubuntu22.04] [--image-name db-gpt]" @@ -19,6 +20,7 @@ usage () { echo " [-i|--pip-index-url pip index url] Pip index url, default: https://pypi.org/simple" echo " [--language en or zh] You language, default: en" echo " [--build-local-code true or false] Whether to use the local project code to package the image, default: false" + echo " [--load-examples true or false] Whether to load examples to default database default: true" echo " [-h|--help] Usage message" } @@ -50,6 +52,11 @@ while [[ $# -gt 0 ]]; do shift shift ;; + --load-examples) + LOAD_EXAMPLES="$2" + shift + shift + ;; -h|--help) help="true" shift @@ -71,5 +78,6 @@ docker build \ --build-arg PIP_INDEX_URL=$PIP_INDEX_URL \ --build-arg LANGUAGE=$LANGUAGE \ --build-arg BUILD_LOCAL_CODE=$BUILD_LOCAL_CODE \ + --build-arg LOAD_EXAMPLES=$LOAD_EXAMPLES \ -f Dockerfile \ -t $IMAGE_NAME $WORK_DIR/../../ diff --git a/docker/base/run_sqlite.sh b/docker/base/run_sqlite.sh new file mode 100755 index 000000000..ed06c2808 --- /dev/null +++ b/docker/base/run_sqlite.sh @@ -0,0 +1,12 @@ +#!/bin/bash + +docker run --gpus all -d \ + -p 5000:5000 \ + -e LOCAL_DB_TYPE=sqlite \ + -e LOCAL_DB_PATH=data/default_sqlite.db \ + -e LLM_MODEL=vicuna-13b-v1.5 \ + -e LANGUAGE=zh \ + -v /data:/data \ + -v /data/models:/app/models \ + --name dbgpt \ + eosphorosai/dbgpt \ No newline at end of file diff --git a/docker/base/run_sqlite_proxyllm.sh b/docker/base/run_sqlite_proxyllm.sh new file mode 100755 index 000000000..a707953cc --- /dev/null +++ b/docker/base/run_sqlite_proxyllm.sh @@ -0,0 +1,18 @@ +#!/bin/bash + +# Your api key +PROXY_API_KEY="$PROXY_API_KEY" +PROXY_SERVER_URL="${PROXY_SERVER_URL-'https://api.openai.com/v1/chat/completions'}" + +docker run --gpus all -d \ + -p 5000:5000 \ + -e LOCAL_DB_TYPE=sqlite \ + -e LOCAL_DB_PATH=data/default_sqlite.db \ + -e LLM_MODEL=proxyllm \ + -e PROXY_API_KEY=$PROXY_API_KEY \ + -e PROXY_SERVER_URL=$PROXY_SERVER_URL \ + -e LANGUAGE=zh \ + -v /data:/data \ + -v /data/models:/app/models \ + --name dbgpt \ + eosphorosai/dbgpt \ No newline at end of file diff --git a/docker/examples/sqls/case_1_student_manager.sql b/docker/examples/sqls/case_1_student_manager_mysql.sql similarity index 100% rename from docker/examples/sqls/case_1_student_manager.sql rename to docker/examples/sqls/case_1_student_manager_mysql.sql diff --git a/docker/examples/sqls/case_1_student_manager_sqlite.sql b/docker/examples/sqls/case_1_student_manager_sqlite.sql new file mode 100644 index 000000000..0f2807c84 --- /dev/null +++ b/docker/examples/sqls/case_1_student_manager_sqlite.sql @@ -0,0 +1,59 @@ +CREATE TABLE students ( + student_id INTEGER PRIMARY KEY, + student_name VARCHAR(100), + major VARCHAR(100), + year_of_enrollment INTEGER, + student_age INTEGER +); + +CREATE TABLE courses ( + course_id INTEGER PRIMARY KEY, + course_name VARCHAR(100), + credit REAL +); + +CREATE TABLE scores ( + student_id INTEGER, + course_id INTEGER, + score INTEGER, + semester VARCHAR(50), + PRIMARY KEY (student_id, course_id), + FOREIGN KEY (student_id) REFERENCES students(student_id), + FOREIGN KEY (course_id) REFERENCES courses(course_id) +); + +INSERT INTO students (student_id, student_name, major, year_of_enrollment, student_age) VALUES +(1, '张三', '计算机科学', 2020, 20), +(2, '李四', '计算机科学', 2021, 19), +(3, '王五', '物理学', 2020, 21), +(4, '赵六', '数学', 2021, 19), +(5, '周七', '计算机科学', 2022, 18), +(6, '吴八', '物理学', 2020, 21), +(7, '郑九', '数学', 2021, 19), +(8, '孙十', '计算机科学', 2022, 18), +(9, '刘十一', '物理学', 2020, 21), +(10, '陈十二', '数学', 2021, 19); + +INSERT INTO courses (course_id, course_name, credit) VALUES +(1, '计算机基础', 3), +(2, '数据结构', 4), +(3, '高等物理', 3), +(4, '线性代数', 4), +(5, '微积分', 5), +(6, '编程语言', 4), +(7, '量子力学', 3), +(8, '概率论', 4), +(9, '数据库系统', 4), +(10, '计算机网络', 4); + +INSERT INTO scores (student_id, course_id, score, semester) VALUES +(1, 1, 90, '2020年秋季'), +(1, 2, 85, '2021年春季'), +(2, 1, 88, '2021年秋季'), +(2, 2, 90, '2022年春季'), +(3, 3, 92, '2020年秋季'), +(3, 4, 85, '2021年春季'), +(4, 3, 88, '2021年秋季'), +(4, 4, 86, '2022年春季'), +(5, 1, 90, '2022年秋季'), +(5, 2, 87, '2023年春季'); diff --git a/docker/examples/sqls/case_2_ecom.sql b/docker/examples/sqls/case_2_ecom_mysql.sql similarity index 100% rename from docker/examples/sqls/case_2_ecom.sql rename to docker/examples/sqls/case_2_ecom_mysql.sql diff --git a/docker/examples/sqls/case_2_ecom_sqlite.sql b/docker/examples/sqls/case_2_ecom_sqlite.sql new file mode 100644 index 000000000..d8ecdd94e --- /dev/null +++ b/docker/examples/sqls/case_2_ecom_sqlite.sql @@ -0,0 +1,59 @@ +CREATE TABLE users ( + user_id INTEGER PRIMARY KEY, + user_name VARCHAR(100), + user_email VARCHAR(100), + registration_date DATE, + user_country VARCHAR(100) +); + +CREATE TABLE products ( + product_id INTEGER PRIMARY KEY, + product_name VARCHAR(100), + product_price REAL +); + +CREATE TABLE orders ( + order_id INTEGER PRIMARY KEY, + user_id INTEGER, + product_id INTEGER, + quantity INTEGER, + order_date DATE, + FOREIGN KEY (user_id) REFERENCES users(user_id), + FOREIGN KEY (product_id) REFERENCES products(product_id) +); + +INSERT INTO users (user_id, user_name, user_email, registration_date, user_country) VALUES +(1, 'John', 'john@gmail.com', '2020-01-01', 'USA'), +(2, 'Mary', 'mary@gmail.com', '2021-01-01', 'UK'), +(3, 'Bob', 'bob@gmail.com', '2020-01-01', 'USA'), +(4, 'Alice', 'alice@gmail.com', '2021-01-01', 'UK'), +(5, 'Charlie', 'charlie@gmail.com', '2020-01-01', 'USA'), +(6, 'David', 'david@gmail.com', '2021-01-01', 'UK'), +(7, 'Eve', 'eve@gmail.com', '2020-01-01', 'USA'), +(8, 'Frank', 'frank@gmail.com', '2021-01-01', 'UK'), +(9, 'Grace', 'grace@gmail.com', '2020-01-01', 'USA'), +(10, 'Helen', 'helen@gmail.com', '2021-01-01', 'UK'); + +INSERT INTO products (product_id, product_name, product_price) VALUES +(1, 'iPhone', 699), +(2, 'Samsung Galaxy', 599), +(3, 'iPad', 329), +(4, 'Macbook', 1299), +(5, 'Apple Watch', 399), +(6, 'AirPods', 159), +(7, 'Echo', 99), +(8, 'Kindle', 89), +(9, 'Fire TV Stick', 39), +(10, 'Echo Dot', 49); + +INSERT INTO orders (order_id, user_id, product_id, quantity, order_date) VALUES +(1, 1, 1, 1, '2022-01-01'), +(2, 1, 2, 1, '2022-02-01'), +(3, 2, 3, 2, '2022-03-01'), +(4, 2, 4, 1, '2022-04-01'), +(5, 3, 5, 2, '2022-05-01'), +(6, 3, 6, 3, '2022-06-01'), +(7, 4, 7, 2, '2022-07-01'), +(8, 4, 8, 1, '2022-08-01'), +(9, 5, 9, 2, '2022-09-01'), +(10, 5, 10, 3, '2022-10-01'); diff --git a/docker/examples/sqls/test_case_info.sql b/docker/examples/sqls/test_case_info_mysql.sql similarity index 100% rename from docker/examples/sqls/test_case_info.sql rename to docker/examples/sqls/test_case_info_mysql.sql diff --git a/docker/examples/sqls/test_case_info_sqlite.sql b/docker/examples/sqls/test_case_info_sqlite.sql new file mode 100644 index 000000000..433e57f6f --- /dev/null +++ b/docker/examples/sqls/test_case_info_sqlite.sql @@ -0,0 +1,17 @@ +CREATE TABLE test_cases ( + case_id INTEGER PRIMARY KEY AUTOINCREMENT, + scenario_name VARCHAR(100), + scenario_description TEXT, + test_question VARCHAR(500), + expected_sql TEXT, + correct_output TEXT +); + + +INSERT INTO test_cases (scenario_name, scenario_description, test_question, expected_sql, correct_output) VALUES +('学校管理系统', '测试SQL助手的联合查询，条件查询和排序功能', '查询所有学生的姓名，专业和成绩，按成绩降序排序', 'SELECT students.student_name, students.major, scores.score FROM students JOIN scores ON students.student_id = scores.student_id ORDER BY scores.score DESC;', '返回所有学生的姓名，专业和成绩，按成绩降序排序的结果'), +('学校管理系统', '测试SQL助手的联合查询，条件查询和排序功能', '查询计算机科学专业的学生的平均成绩', 'SELECT AVG(scores.score) as avg_score FROM students JOIN scores ON students.student_id = scores.student_id WHERE students.major = ''计算机科学'';', '返回计算机科学专业学生的平均成绩'), +('学校管理系统', '测试SQL助手的联合查询，条件查询和排序功能', '查询哪些学生在2023年秋季学期的课程学分总和超过15', 'SELECT students.student_name FROM students JOIN scores ON students.student_id = scores.student_id JOIN courses ON scores.course_id = courses.course_id WHERE scores.semester = ''2023年秋季'' GROUP BY students.student_id HAVING SUM(courses.credit) > 15;', '返回在2023年秋季学期的课程学分总和超过15的学生的姓名'), +('电商系统', '测试SQL助手的数据聚合和分组功能', '查询每个用户的总订单数量', 'SELECT users.user_name, COUNT(orders.order_id) as order_count FROM users JOIN orders ON users.user_id = orders.user_id GROUP BY users.user_id;', '返回每个用户的总订单数量'), +('电商系统', '测试SQL助手的数据聚合和分组功能', '查询每种商品的总销售额', 'SELECT products.product_name, SUM(products.product_price * orders.quantity) as total_sales FROM products JOIN orders ON products.product_id = orders.product_id GROUP BY products.product_id;', '返回每种商品的总销售额'), +('电商系统', '测试SQL助手的数据聚合和分组功能', '查询2023年最受欢迎的商品（订单数量最多的商品）', 'SELECT products.product_name FROM products JOIN orders ON products.product_id = orders.product_id WHERE YEAR(orders.order_date) = 2023 GROUP BY products.product_id ORDER BY COUNT(orders.order_id) DESC LIMIT 1;', '返回2023年最受欢迎的商品（订单数量最多的商品）的名称'); diff --git a/docker/examples/sqls/user_config.sql b/docker/examples/sqls/user_config_mysql.sql similarity index 100% rename from docker/examples/sqls/user_config.sql rename to docker/examples/sqls/user_config_mysql.sql diff --git a/pilot/configs/config.py b/pilot/configs/config.py index 609d946c2..2a341d017 100644 --- a/pilot/configs/config.py +++ b/pilot/configs/config.py @@ -123,7 +123,7 @@ class Config(metaclass=Singleton): ### default Local database connection configuration self.LOCAL_DB_HOST = os.getenv("LOCAL_DB_HOST") self.LOCAL_DB_PATH = os.getenv("LOCAL_DB_PATH", "") - self.LOCAL_DB_TYPE = os.getenv("LOCAL_DB_TYPE") + self.LOCAL_DB_TYPE = os.getenv("LOCAL_DB_TYPE", "mysql") if self.LOCAL_DB_HOST is None and self.LOCAL_DB_PATH == "": self.LOCAL_DB_HOST = "127.0.0.1" diff --git a/pilot/connections/__init__.py b/pilot/connections/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pilot/connections/manages/connection_manager.py b/pilot/connections/manages/connection_manager.py index 19b9e64cf..6dc1a0222 100644 --- a/pilot/connections/manages/connection_manager.py +++ b/pilot/connections/manages/connection_manager.py @@ -93,19 +93,26 @@ class ConnectManager: db_name = CFG.LOCAL_DB_NAME db_type = CFG.LOCAL_DB_TYPE db_path = CFG.LOCAL_DB_PATH + if not db_type: + # Default file database type + db_type = DBType.DuckDb.value() if not db_name: - if db_type is None or db_type == DBType.DuckDb.value(): - # file db is duckdb - db_name = self.storage.get_file_db_name(db_path) - db_type = DBType.DuckDb.value() - else: - db_name = DBType.parse_file_db_name_from_path(db_type, db_path) + db_type, db_name = self._parse_file_db_info(db_type, db_path) if db_name: print( f"Add file db, db_name: {db_name}, db_type: {db_type}, db_path: {db_path}" ) self.storage.add_file_db(db_name, db_type, db_path) + def _parse_file_db_info(self, db_type: str, db_path: str): + if db_type is None or db_type == DBType.DuckDb.value(): + # file db is duckdb + db_name = self.storage.get_file_db_name(db_path) + db_type = DBType.DuckDb.value() + else: + db_name = DBType.parse_file_db_name_from_path(db_type, db_path) + return db_type, db_name + def get_connect(self, db_name): db_config = self.storage.get_db_config(db_name) db_type = DBType.of_db_type(db_config.get("db_type")) diff --git a/pilot/connections/rdbms/base_dao.py b/pilot/connections/rdbms/base_dao.py new file mode 100644 index 000000000..82edb404f --- /dev/null +++ b/pilot/connections/rdbms/base_dao.py @@ -0,0 +1,83 @@ +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker +from pilot.configs.config import Config +from pilot.common.schema import DBType +from pilot.connections.rdbms.base import RDBMSDatabase +from pilot.logs import logger + +CFG = Config() + + +class BaseDao: + def __init__( + self, orm_base=None, database: str = None, create_not_exist_table: bool = False + ) -> None: + """BaseDAO, If the current database is a file database and create_not_exist_table=True, we will automatically create a table that does not exist""" + self._orm_base = orm_base + self._database = database + self._create_not_exist_table = create_not_exist_table + + self._db_engine = None + self._session = None + self._connection = None + + @property + def db_engine(self): + if not self._db_engine: + # lazy loading + db_engine, connection = _get_db_engine( + self._orm_base, self._database, self._create_not_exist_table + ) + self._db_engine = db_engine + self._connection = connection + return self._db_engine + + @property + def Session(self): + if not self._session: + self._session = sessionmaker(bind=self.db_engine) + return self._session + + +def _get_db_engine( + orm_base=None, database: str = None, create_not_exist_table: bool = False +): + db_engine = None + connection: RDBMSDatabase = None + + db_type = DBType.of_db_type(CFG.LOCAL_DB_TYPE) + if db_type is None or db_type == DBType.Mysql: + # default database + db_engine = create_engine( + f"mysql+pymysql://{CFG.LOCAL_DB_USER}:{CFG.LOCAL_DB_PASSWORD}@{CFG.LOCAL_DB_HOST}:{CFG.LOCAL_DB_PORT}/{database}", + echo=True, + ) + else: + db_namager = CFG.LOCAL_DB_MANAGE + if not db_namager: + raise Exception( + "LOCAL_DB_MANAGE is not initialized, please check the system configuration" + ) + if db_type.is_file_db(): + db_path = CFG.LOCAL_DB_PATH + if db_path is None or db_path == "": + raise ValueError( + "You LOCAL_DB_TYPE is file db, but LOCAL_DB_PATH is not configured, please configure LOCAL_DB_PATH in you .env file" + ) + _, database = db_namager._parse_file_db_info(db_type.value(), db_path) + logger.info( + f"Current DAO database is file database, db_type: {db_type.value()}, db_path: {db_path}, db_name: {database}" + ) + logger.info(f"Get DAO database connection with database name {database}") + connection: RDBMSDatabase = db_namager.get_connect(database) + if not isinstance(connection, RDBMSDatabase): + raise ValueError( + "Currently only supports `RDBMSDatabase` database as the underlying database of BaseDao, please check your database configuration" + ) + db_engine = connection._engine + + if db_type.is_file_db() and orm_base is not None and create_not_exist_table: + logger.info("Current database is file database, create not exist table") + orm_base.metadata.create_all(db_engine) + + return db_engine, connection diff --git a/pilot/connections/rdbms/conn_sqlite.py b/pilot/connections/rdbms/conn_sqlite.py index f4c8084a2..339af025a 100644 --- a/pilot/connections/rdbms/conn_sqlite.py +++ b/pilot/connections/rdbms/conn_sqlite.py @@ -22,6 +22,8 @@ class SQLiteConnect(RDBMSDatabase): ) -> RDBMSDatabase: """Construct a SQLAlchemy engine from URI.""" _engine_args = engine_args or {} + _engine_args["connect_args"] = {"check_same_thread": False} + # _engine_args["echo"] = True return cls(create_engine("sqlite:///" + file_path, **_engine_args), **kwargs) def get_indexes(self, table_name): diff --git a/pilot/server/knowledge/api.py b/pilot/server/knowledge/api.py index b6aef1ab3..51ee7f924 100644 --- a/pilot/server/knowledge/api.py +++ b/pilot/server/knowledge/api.py @@ -144,10 +144,8 @@ async def document_upload( request = KnowledgeDocumentRequest() request.doc_name = doc_name request.doc_type = doc_type - request.content = ( - os.path.join( - KNOWLEDGE_UPLOAD_ROOT_PATH, space_name, doc_file.filename - ), + request.content = os.path.join( + KNOWLEDGE_UPLOAD_ROOT_PATH, space_name, doc_file.filename ) return Result.succ( knowledge_space_service.create_knowledge_document( diff --git a/pilot/server/knowledge/chunk_db.py b/pilot/server/knowledge/chunk_db.py index d3b58fc91..7ca9bd461 100644 --- a/pilot/server/knowledge/chunk_db.py +++ b/pilot/server/knowledge/chunk_db.py @@ -5,7 +5,7 @@ from sqlalchemy import Column, String, DateTime, Integer, Text, create_engine, f from sqlalchemy.orm import declarative_base, sessionmaker from pilot.configs.config import Config - +from pilot.connections.rdbms.base_dao import BaseDao CFG = Config() @@ -27,14 +27,11 @@ class DocumentChunkEntity(Base): return f"DocumentChunkEntity(id={self.id}, doc_name='{self.doc_name}', doc_type='{self.doc_type}', document_id='{self.document_id}', content='{self.content}', meta_info='{self.meta_info}', gmt_created='{self.gmt_created}', gmt_modified='{self.gmt_modified}')" -class DocumentChunkDao: +class DocumentChunkDao(BaseDao): def __init__(self): - database = "knowledge_management" - self.db_engine = create_engine( - f"mysql+pymysql://{CFG.LOCAL_DB_USER}:{CFG.LOCAL_DB_PASSWORD}@{CFG.LOCAL_DB_HOST}:{CFG.LOCAL_DB_PORT}/{database}", - echo=True, + super().__init__( + database="knowledge_management", orm_base=Base, create_not_exist_table=True ) - self.Session = sessionmaker(bind=self.db_engine) def create_documents_chunks(self, documents: List): session = self.Session() diff --git a/pilot/server/knowledge/document_db.py b/pilot/server/knowledge/document_db.py index 8b01c7877..1a6c4f8ac 100644 --- a/pilot/server/knowledge/document_db.py +++ b/pilot/server/knowledge/document_db.py @@ -4,7 +4,7 @@ from sqlalchemy import Column, String, DateTime, Integer, Text, create_engine, f from sqlalchemy.orm import declarative_base, sessionmaker from pilot.configs.config import Config - +from pilot.connections.rdbms.base_dao import BaseDao CFG = Config() @@ -19,7 +19,7 @@ class KnowledgeDocumentEntity(Base): space = Column(String(100)) chunk_size = Column(Integer) status = Column(String(100)) - last_sync = Column(String(100)) + last_sync = Column(DateTime) content = Column(Text) result = Column(Text) vector_ids = Column(Text) @@ -30,14 +30,11 @@ class KnowledgeDocumentEntity(Base): return f"KnowledgeDocumentEntity(id={self.id}, doc_name='{self.doc_name}', doc_type='{self.doc_type}', chunk_size='{self.chunk_size}', status='{self.status}', last_sync='{self.last_sync}', content='{self.content}', result='{self.result}', gmt_created='{self.gmt_created}', gmt_modified='{self.gmt_modified}')" -class KnowledgeDocumentDao: +class KnowledgeDocumentDao(BaseDao): def __init__(self): - database = "knowledge_management" - self.db_engine = create_engine( - f"mysql+pymysql://{CFG.LOCAL_DB_USER}:{CFG.LOCAL_DB_PASSWORD}@{CFG.LOCAL_DB_HOST}:{CFG.LOCAL_DB_PORT}/{database}", - echo=True, + super().__init__( + database="knowledge_management", orm_base=Base, create_not_exist_table=True ) - self.Session = sessionmaker(bind=self.db_engine) def create_knowledge_document(self, document: KnowledgeDocumentEntity): session = self.Session() diff --git a/pilot/server/knowledge/space_db.py b/pilot/server/knowledge/space_db.py index 57a909e06..cb2640b00 100644 --- a/pilot/server/knowledge/space_db.py +++ b/pilot/server/knowledge/space_db.py @@ -2,11 +2,11 @@ from datetime import datetime from sqlalchemy import Column, Integer, Text, String, DateTime, create_engine from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy.orm import sessionmaker from pilot.configs.config import Config - from pilot.server.knowledge.request.request import KnowledgeSpaceRequest -from sqlalchemy.orm import sessionmaker +from pilot.connections.rdbms.base_dao import BaseDao CFG = Config() Base = declarative_base() @@ -27,14 +27,11 @@ class KnowledgeSpaceEntity(Base): return f"KnowledgeSpaceEntity(id={self.id}, name='{self.name}', vector_type='{self.vector_type}', desc='{self.desc}', owner='{self.owner}' context='{self.context}', gmt_created='{self.gmt_created}', gmt_modified='{self.gmt_modified}')" -class KnowledgeSpaceDao: +class KnowledgeSpaceDao(BaseDao): def __init__(self): - database = "knowledge_management" - self.db_engine = create_engine( - f"mysql+pymysql://{CFG.LOCAL_DB_USER}:{CFG.LOCAL_DB_PASSWORD}@{CFG.LOCAL_DB_HOST}:{CFG.LOCAL_DB_PORT}/{database}", - echo=True, + super().__init__( + database="knowledge_management", orm_base=Base, create_not_exist_table=True ) - self.Session = sessionmaker(bind=self.db_engine) def create_knowledge_space(self, space: KnowledgeSpaceRequest): session = self.Session()