VectorCode/src/vectorcode/main.py at db7e74c47d396afd55b676783c569b0ffa2b4a57 · Davidyz/VectorCode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
import asyncio
import logging
import os
import sys
import traceback

import httpx

from vectorcode import __version__
from vectorcode.cli_utils import (
    CliAction,
    config_logging,
    find_project_root,
    load_config_file,
    parse_cli_args,
)
from vectorcode.common import ClientManager

logger = logging.getLogger(name=__name__)


async def async_main():
    cli_args = await parse_cli_args()
    if cli_args.no_stderr:
        sys.stderr = open(os.devnull, "w")

    if cli_args.debug:  # pragma: nocover
        from vectorcode import debugging

        debugging.enable()

    logger.info("Collected CLI arguments: %s", cli_args)

    if cli_args.project_root is None:
        cwd = os.getcwd()
        cli_args.project_root = (
            find_project_root(cwd, ".vectorcode")
            or find_project_root(cwd, ".git")
            or cwd
        )

    logger.info(f"Project root is set to {cli_args.project_root}")

    try:
        final_configs = await (
            await load_config_file(cli_args.project_root)
        ).merge_from(cli_args)
    except IOError as e:
        traceback.print_exception(e, file=sys.stderr)
        return 1

    logger.info("Final configuration has been built: %s", final_configs)

    match cli_args.action:
        case CliAction.check:
            from vectorcode.subcommands import check

            return await check(cli_args)
        case CliAction.init:
            from vectorcode.subcommands import init

            return await init(cli_args)
        case CliAction.version:
            print(__version__)
            return 0
        case CliAction.prompts:
            from vectorcode.subcommands import prompts

            return prompts(cli_args)
        case CliAction.chunks:
            from vectorcode.subcommands import chunks

            return await chunks(final_configs)

    if final_configs.pipe:  # pragma: nocover
        # NOTE: NNCF (intel GPU acceleration for sentence transformer) keeps showing logs.
        # This disables logs below ERROR so that it doesn't hurt the `pipe` output.
        logging.disable(logging.ERROR - 1)

    return_val = 0
    try:
        match final_configs.action:
            case CliAction.query:
                from vectorcode.subcommands import query

                return_val = await query(final_configs)
            case CliAction.vectorise:
                from vectorcode.subcommands import vectorise

                return_val = await vectorise(final_configs)
            case CliAction.drop:
                from vectorcode.subcommands import drop

                return_val = await drop(final_configs)
            case CliAction.ls:
                from vectorcode.subcommands import ls

                return_val = await ls(final_configs)
            case CliAction.update:
                from vectorcode.subcommands import update

                return_val = await update(final_configs)
            case CliAction.clean:
                from vectorcode.subcommands import clean

                return_val = await clean(final_configs)
            case CliAction.files:
                from vectorcode.subcommands import files

                return_val = await files(final_configs)
    except Exception as e:
        return_val = 1
        if isinstance(e, httpx.RemoteProtocolError):  # pragma: nocover
            e.add_note(
                f"Please verify that {final_configs.db_url} is a working chromadb server."
            )
        logger.error(traceback.format_exc())
    finally:
        await ClientManager().kill_servers()
        return return_val


def main():  # pragma: nocover
    config_logging("vectorcode")
    return asyncio.run(async_main())


if __name__ == "__main__":  # pragma: nocover
    sys.exit(main())