mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-09-29 21:40:44 +00:00
122 lines
2.2 KiB
Python
122 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
DB-GPT command line tools.
|
|
|
|
You can use it for some background management:
|
|
- Lots of knowledge document initialization.
|
|
- Load the data into the database.
|
|
- Show server status
|
|
- ...
|
|
|
|
|
|
Maybe move this to pilot module and append to console_scripts in the future.
|
|
|
|
"""
|
|
import sys
|
|
import click
|
|
import os
|
|
|
|
sys.path.append(
|
|
os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
|
|
)
|
|
|
|
|
|
from pilot.configs.model_config import DATASETS_DIR
|
|
|
|
from tools.cli.knowledge_client import knowledge_init
|
|
|
|
API_ADDRESS: str = "http://127.0.0.1:5000"
|
|
|
|
|
|
@click.group()
|
|
@click.option(
|
|
"--api_address",
|
|
required=False,
|
|
default="http://127.0.0.1:5000",
|
|
type=str,
|
|
help="Api server address",
|
|
)
|
|
@click.version_option()
|
|
def cli(api_address: str):
|
|
global API_ADDRESS
|
|
API_ADDRESS = api_address
|
|
|
|
|
|
@cli.command()
|
|
@click.option(
|
|
"--vector_name",
|
|
required=False,
|
|
type=str,
|
|
default="default",
|
|
help="Your vector store name",
|
|
)
|
|
@click.option(
|
|
"--vector_store_type",
|
|
required=False,
|
|
type=str,
|
|
default="Chroma",
|
|
help="Vector store type",
|
|
)
|
|
@click.option(
|
|
"--local_doc_dir",
|
|
required=False,
|
|
type=str,
|
|
default=DATASETS_DIR,
|
|
help="Your document directory",
|
|
)
|
|
@click.option(
|
|
"--skip_wrong_doc",
|
|
required=False,
|
|
type=bool,
|
|
default=False,
|
|
help="Skip wrong document",
|
|
)
|
|
@click.option(
|
|
"--max_workers",
|
|
required=False,
|
|
type=int,
|
|
default=None,
|
|
help="The maximum number of threads that can be used to upload document",
|
|
)
|
|
@click.option(
|
|
"-v",
|
|
"--verbose",
|
|
required=False,
|
|
is_flag=True,
|
|
hidden=True,
|
|
help="Show debuggging information.",
|
|
)
|
|
def knowledge(
|
|
vector_name: str,
|
|
vector_store_type: str,
|
|
local_doc_dir: str,
|
|
skip_wrong_doc: bool,
|
|
max_workers: int,
|
|
verbose: bool,
|
|
):
|
|
"""Knowledge command line tool"""
|
|
knowledge_init(
|
|
API_ADDRESS,
|
|
vector_name,
|
|
vector_store_type,
|
|
local_doc_dir,
|
|
skip_wrong_doc,
|
|
verbose,
|
|
max_workers,
|
|
)
|
|
|
|
|
|
# knowledge command
|
|
cli.add_command(knowledge)
|
|
# TODO add more command
|
|
|
|
|
|
def main():
|
|
return cli()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|