diff --git a/Makefile b/Makefile index c8b06e756..cddfa307b 100644 --- a/Makefile +++ b/Makefile @@ -81,6 +81,19 @@ clean: ## Clean up the environment find . -type d -name '.pytest_cache' -delete find . -type d -name '.coverage' -delete +.PHONY: clean-dist +clean-dist: ## Clean up the distribution + rm -rf dist/ *.egg-info build/ + +.PHONY: package +package: clean-dist ## Package the project for distribution + IS_DEV_MODE=false python setup.py sdist bdist_wheel + +.PHONY: upload +upload: package ## Upload the package to PyPI + # upload to testpypi: twine upload --repository testpypi dist/* + twine upload dist/* + .PHONY: help help: ## Display this help screen @echo "Available commands:" diff --git a/dbgpt/__init__.py b/dbgpt/__init__.py index 1660c5ddc..14efc6ca3 100644 --- a/dbgpt/__init__.py +++ b/dbgpt/__init__.py @@ -1,12 +1,18 @@ -from dbgpt.component import BaseComponent, SystemApp - -__ALL__ = ["SystemApp", "BaseComponent"] +"""DB-GPT: Next Generation Data Interaction Solution with LLMs. +""" +from dbgpt import _version # noqa: E402 +from dbgpt.component import BaseComponent, SystemApp # noqa: F401 _CORE_LIBS = ["core", "rag", "model", "agent", "datasource", "vis", "storage", "train"] _SERVE_LIBS = ["serve"] _LIBS = _CORE_LIBS + _SERVE_LIBS +__version__ = _version.version + +__ALL__ = ["__version__", "SystemApp", "BaseComponent"] + + def __getattr__(name: str): # Lazy load import importlib diff --git a/dbgpt/_version.py b/dbgpt/_version.py new file mode 100644 index 000000000..e7cc0113f --- /dev/null +++ b/dbgpt/_version.py @@ -0,0 +1 @@ +version = "0.4.7" diff --git a/dbgpt/agent/agents/expand/retrieve_summary_assistant_agent.py b/dbgpt/agent/agents/expand/retrieve_summary_assistant_agent.py index 8d59fe39f..21783f3d7 100644 --- a/dbgpt/agent/agents/expand/retrieve_summary_assistant_agent.py +++ b/dbgpt/agent/agents/expand/retrieve_summary_assistant_agent.py @@ -579,7 +579,7 @@ def _num_token_from_text(self, text: str, model: str = "gpt-3.5-turbo-0613"): from dbgpt.agent.agents.agent import AgentContext from dbgpt.agent.agents.user_proxy_agent import UserProxyAgent from dbgpt.core.interface.llm import ModelMetadata - from dbgpt.model import OpenAILLMClient + from dbgpt.model.proxy import OpenAILLMClient llm_client = OpenAILLMClient() context: AgentContext = AgentContext( diff --git a/dbgpt/app/chat_adapter.py b/dbgpt/app/chat_adapter.py index d4eb3f3a7..0d537d479 100644 --- a/dbgpt/app/chat_adapter.py +++ b/dbgpt/app/chat_adapter.py @@ -9,7 +9,7 @@ from typing import Dict, List, Tuple from dbgpt.core.interface.message import ModelMessage, ModelMessageRoleType -from dbgpt.model.conversation import Conversation, get_conv_template +from dbgpt.model.llm.conversation import Conversation, get_conv_template class BaseChatAdpter: @@ -21,7 +21,7 @@ def match(self, model_path: str): def get_generate_stream_func(self, model_path: str): """Return the generate stream handler func""" - from dbgpt.model.inference import generate_stream + from dbgpt.model.llm.inference import generate_stream return generate_stream diff --git a/dbgpt/app/scene/base_chat.py b/dbgpt/app/scene/base_chat.py index f48feceb0..10f738d3b 100644 --- a/dbgpt/app/scene/base_chat.py +++ b/dbgpt/app/scene/base_chat.py @@ -171,13 +171,13 @@ def llm_client(self) -> LLMClient: async def call_llm_operator(self, request: ModelRequest) -> ModelOutput: llm_task = build_cached_chat_operator(self.llm_client, False, CFG.SYSTEM_APP) - return await llm_task.call(call_data={"data": request}) + return await llm_task.call(call_data=request) async def call_streaming_operator( self, request: ModelRequest ) -> AsyncIterator[ModelOutput]: llm_task = build_cached_chat_operator(self.llm_client, True, CFG.SYSTEM_APP) - async for out in await llm_task.call_stream(call_data={"data": request}): + async for out in await llm_task.call_stream(call_data=request): yield out def do_action(self, prompt_response): @@ -251,11 +251,9 @@ async def _build_model_request(self) -> ModelRequest: str_history=self.prompt_template.str_history, request_context=req_ctx, ) - node_input = { - "data": ChatComposerInput( - messages=self.history_messages, prompt_dict=input_values - ) - } + node_input = ChatComposerInput( + messages=self.history_messages, prompt_dict=input_values + ) # llm_messages = self.generate_llm_messages() model_request: ModelRequest = await node.call(call_data=node_input) model_request.context.cache_enable = self.model_cache_enable diff --git a/dbgpt/app/scene/operator/app_operator.py b/dbgpt/app/scene/operator/app_operator.py index c9e7f03b6..91f65e357 100644 --- a/dbgpt/app/scene/operator/app_operator.py +++ b/dbgpt/app/scene/operator/app_operator.py @@ -87,7 +87,7 @@ async def map(self, input_value: ChatComposerInput) -> ModelRequest: end_node: BaseOperator = self._sub_compose_dag.leaf_nodes[0] # Sub dag, use the same dag context in the parent dag messages = await end_node.call( - call_data={"data": input_value}, dag_ctx=self.current_dag_context + call_data=input_value, dag_ctx=self.current_dag_context ) span_id = self._request_context.span_id model_request = ModelRequest.build_request( diff --git a/dbgpt/component.py b/dbgpt/component.py index 2468906b9..3b26c381b 100644 --- a/dbgpt/component.py +++ b/dbgpt/component.py @@ -1,3 +1,7 @@ +"""Component module for dbgpt. + +Manages the lifecycle and registration of components. +""" from __future__ import annotations import asyncio diff --git a/dbgpt/core/awel/__init__.py b/dbgpt/core/awel/__init__.py index 84e2fc775..c97ef1186 100644 --- a/dbgpt/core/awel/__init__.py +++ b/dbgpt/core/awel/__init__.py @@ -22,6 +22,7 @@ JoinOperator, MapOperator, ReduceStreamOperator, + TriggerOperator, ) from .operator.stream_operator import ( StreamifyAbsOperator, @@ -50,6 +51,7 @@ "BaseOperator", "JoinOperator", "ReduceStreamOperator", + "TriggerOperator", "MapOperator", "BranchOperator", "InputOperator", @@ -150,4 +152,6 @@ def setup_dev_environment( for trigger in dag.trigger_nodes: trigger_manager.register_trigger(trigger) trigger_manager.after_register() - uvicorn.run(app, host=host, port=port) + if trigger_manager.keep_running(): + # Should keep running + uvicorn.run(app, host=host, port=port) diff --git a/dbgpt/core/awel/operator/base.py b/dbgpt/core/awel/operator/base.py index 0af890e7c..14aa1fb33 100644 --- a/dbgpt/core/awel/operator/base.py +++ b/dbgpt/core/awel/operator/base.py @@ -28,7 +28,7 @@ F = TypeVar("F", bound=FunctionType) -CALL_DATA = Union[Dict, Dict[str, Dict]] +CALL_DATA = Union[Dict[str, Any], Any] class WorkflowRunner(ABC, Generic[T]): @@ -197,6 +197,8 @@ async def call( Returns: OUT: The output of the node after execution. """ + if call_data: + call_data = {"data": call_data} out_ctx = await self._runner.execute_workflow( self, call_data, exist_dag_ctx=dag_ctx ) @@ -242,6 +244,8 @@ async def call_stream( Returns: AsyncIterator[OUT]: An asynchronous iterator over the output stream. """ + if call_data: + call_data = {"data": call_data} out_ctx = await self._runner.execute_workflow( self, call_data, streaming_call=True, exist_dag_ctx=dag_ctx ) diff --git a/dbgpt/core/awel/task/base.py b/dbgpt/core/awel/task/base.py index f0c5712bc..fb816838e 100644 --- a/dbgpt/core/awel/task/base.py +++ b/dbgpt/core/awel/task/base.py @@ -28,6 +28,14 @@ def __bool__(self): SKIP_DATA = _EMPTY_DATA_TYPE() PLACEHOLDER_DATA = _EMPTY_DATA_TYPE() + +def is_empty_data(data: Any): + """Check if the data is empty.""" + if isinstance(data, _EMPTY_DATA_TYPE): + return data in (EMPTY_DATA, SKIP_DATA) + return False + + MapFunc = Union[Callable[[IN], OUT], Callable[[IN], Awaitable[OUT]]] ReduceFunc = Union[Callable[[IN], OUT], Callable[[IN], Awaitable[OUT]]] StreamFunc = Callable[[IN], Awaitable[AsyncIterator[OUT]]] diff --git a/dbgpt/core/awel/task/task_impl.py b/dbgpt/core/awel/task/task_impl.py index 8877c5cfe..4b42443f5 100644 --- a/dbgpt/core/awel/task/task_impl.py +++ b/dbgpt/core/awel/task/task_impl.py @@ -24,7 +24,6 @@ EMPTY_DATA, OUT, PLACEHOLDER_DATA, - SKIP_DATA, InputContext, InputSource, MapFunc, @@ -37,6 +36,7 @@ TaskState, TransformFunc, UnStreamFunc, + is_empty_data, ) logger = logging.getLogger(__name__) @@ -99,7 +99,7 @@ def new_output(self) -> TaskOutput[T]: @property def is_empty(self) -> bool: """Return True if the output data is empty.""" - return self._data == EMPTY_DATA or self._data == SKIP_DATA + return is_empty_data(self._data) @property def is_none(self) -> bool: @@ -171,7 +171,7 @@ def is_stream(self) -> bool: @property def is_empty(self) -> bool: """Return True if the output data is empty.""" - return self._data == EMPTY_DATA or self._data == SKIP_DATA + return is_empty_data(self._data) @property def is_none(self) -> bool: @@ -330,7 +330,7 @@ def _read_data(self, task_ctx: TaskContext) -> Any: """ call_data = task_ctx.call_data data = call_data.get("data", EMPTY_DATA) if call_data else EMPTY_DATA - if data == EMPTY_DATA: + if is_empty_data(data): raise ValueError("No call data for current SimpleCallDataInputSource") return data diff --git a/dbgpt/core/awel/trigger/http_trigger.py b/dbgpt/core/awel/trigger/http_trigger.py index a48dfe385..c77bd8b7e 100644 --- a/dbgpt/core/awel/trigger/http_trigger.py +++ b/dbgpt/core/awel/trigger/http_trigger.py @@ -1,12 +1,8 @@ """Http trigger for AWEL.""" -from __future__ import annotations - import logging from enum import Enum from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union, cast -from starlette.requests import Request - from dbgpt._private.pydantic import BaseModel from ..dag.base import DAG @@ -15,9 +11,10 @@ if TYPE_CHECKING: from fastapi import APIRouter + from starlette.requests import Request -RequestBody = Union[Type[Request], Type[BaseModel], Type[str]] -StreamingPredictFunc = Callable[[Union[Request, BaseModel]], bool] + RequestBody = Union[Type[Request], Type[BaseModel], Type[str]] + StreamingPredictFunc = Callable[[Union[Request, BaseModel]], bool] logger = logging.getLogger(__name__) @@ -32,9 +29,9 @@ def __init__( self, endpoint: str, methods: Optional[Union[str, List[str]]] = "GET", - request_body: Optional[RequestBody] = None, + request_body: Optional["RequestBody"] = None, streaming_response: bool = False, - streaming_predict_func: Optional[StreamingPredictFunc] = None, + streaming_predict_func: Optional["StreamingPredictFunc"] = None, response_model: Optional[Type] = None, response_headers: Optional[Dict[str, str]] = None, response_media_type: Optional[str] = None, @@ -69,6 +66,7 @@ def mount_to_router(self, router: "APIRouter") -> None: router (APIRouter): The router to mount the trigger. """ from fastapi import Depends + from starlette.requests import Request methods = [self._methods] if isinstance(self._methods, str) else self._methods @@ -114,8 +112,10 @@ async def route_function(body=Depends(_request_body_dependency)): async def _parse_request_body( - request: Request, request_body_cls: Optional[RequestBody] + request: "Request", request_body_cls: Optional["RequestBody"] ): + from starlette.requests import Request + if not request_body_cls: return None if request_body_cls == Request: @@ -152,7 +152,7 @@ async def _trigger_dag( raise ValueError("HttpTrigger just support one leaf node in dag") end_node = cast(BaseOperator, leaf_nodes[0]) if not streaming_response: - return await end_node.call(call_data={"data": body}) + return await end_node.call(call_data=body) else: headers = response_headers media_type = response_media_type if response_media_type else "text/event-stream" @@ -163,7 +163,7 @@ async def _trigger_dag( "Connection": "keep-alive", "Transfer-Encoding": "chunked", } - generator = await end_node.call_stream(call_data={"data": body}) + generator = await end_node.call_stream(call_data=body) background_tasks = BackgroundTasks() background_tasks.add_task(dag._after_dag_end) return StreamingResponse( diff --git a/dbgpt/core/awel/trigger/trigger_manager.py b/dbgpt/core/awel/trigger/trigger_manager.py index c9baed58d..608e68fa9 100644 --- a/dbgpt/core/awel/trigger/trigger_manager.py +++ b/dbgpt/core/awel/trigger/trigger_manager.py @@ -24,6 +24,14 @@ class TriggerManager(ABC): def register_trigger(self, trigger: Any) -> None: """Register a trigger to current manager.""" + def keep_running(self) -> bool: + """Whether keep running. + + Returns: + bool: Whether keep running, True means keep running, False means stop. + """ + return False + class HttpTriggerManager(TriggerManager): """Http trigger manager. @@ -64,6 +72,8 @@ def register_trigger(self, trigger: Any) -> None: self._trigger_map[trigger_id] = trigger def _init_app(self, system_app: SystemApp): + if not self.keep_running(): + return logger.info( f"Include router {self._router} to prefix path {self._router_prefix}" ) @@ -72,6 +82,14 @@ def _init_app(self, system_app: SystemApp): raise RuntimeError("System app not initialized") app.include_router(self._router, prefix=self._router_prefix, tags=["AWEL"]) + def keep_running(self) -> bool: + """Whether keep running. + + Returns: + bool: Whether keep running, True means keep running, False means stop. + """ + return len(self._trigger_map) > 0 + class DefaultTriggerManager(TriggerManager, BaseComponent): """Default trigger manager for AWEL. @@ -105,3 +123,11 @@ def after_register(self) -> None: """After register, init the trigger manager.""" if self.system_app: self.http_trigger._init_app(self.system_app) + + def keep_running(self) -> bool: + """Whether keep running. + + Returns: + bool: Whether keep running, True means keep running, False means stop. + """ + return self.http_trigger.keep_running() diff --git a/dbgpt/core/interface/operator/composer_operator.py b/dbgpt/core/interface/operator/composer_operator.py index a8896702a..bf7a41204 100644 --- a/dbgpt/core/interface/operator/composer_operator.py +++ b/dbgpt/core/interface/operator/composer_operator.py @@ -70,7 +70,7 @@ async def map(self, input_value: ChatComposerInput) -> ModelRequest: end_node: BaseOperator = cast(BaseOperator, self._sub_compose_dag.leaf_nodes[0]) # Sub dag, use the same dag context in the parent dag return await end_node.call( - call_data={"data": input_value}, dag_ctx=self.current_dag_context + call_data=input_value, dag_ctx=self.current_dag_context ) def _build_composer_dag(self) -> DAG: diff --git a/dbgpt/core/interface/operator/prompt_operator.py b/dbgpt/core/interface/operator/prompt_operator.py index 5049c7026..f31172533 100644 --- a/dbgpt/core/interface/operator/prompt_operator.py +++ b/dbgpt/core/interface/operator/prompt_operator.py @@ -150,7 +150,7 @@ class PromptBuilderOperator( ) ) - single_input = {"data": {"dialect": "mysql"}} + single_input = {"dialect": "mysql"} single_expected_messages = [ ModelMessage( content="Please write a mysql SQL count the length of a field", diff --git a/dbgpt/model/__init__.py b/dbgpt/model/__init__.py index dacfc9a50..0e66b2ff1 100644 --- a/dbgpt/model/__init__.py +++ b/dbgpt/model/__init__.py @@ -1,9 +1,12 @@ -from dbgpt.model.cluster.client import DefaultLLMClient +try: + from dbgpt.model.cluster.client import DefaultLLMClient +except ImportError as exc: + # logging.warning("Can't import dbgpt.model.DefaultLLMClient") + DefaultLLMClient = None -# from dbgpt.model.utils.chatgpt_utils import OpenAILLMClient -from dbgpt.model.proxy.llms.chatgpt import OpenAILLMClient -__ALL__ = [ - "DefaultLLMClient", - "OpenAILLMClient", -] +_exports = [] +if DefaultLLMClient: + _exports.append("DefaultLLMClient") + +__ALL__ = _exports diff --git a/dbgpt/model/loader.py b/dbgpt/model/adapter/loader.py similarity index 99% rename from dbgpt/model/loader.py rename to dbgpt/model/adapter/loader.py index 94872a2cb..6ce4fbe52 100644 --- a/dbgpt/model/loader.py +++ b/dbgpt/model/adapter/loader.py @@ -137,7 +137,7 @@ def loader_with_params( def huggingface_loader(llm_adapter: LLMModelAdapter, model_params: ModelParameters): import torch - from dbgpt.model.compression import compress_module + from dbgpt.model.llm.compression import compress_module device = model_params.device max_memory = None diff --git a/dbgpt/model/adapter/old_adapter.py b/dbgpt/model/adapter/old_adapter.py index e9ed01af4..774ee5ef0 100644 --- a/dbgpt/model/adapter/old_adapter.py +++ b/dbgpt/model/adapter/old_adapter.py @@ -19,7 +19,7 @@ from dbgpt.model.adapter.base import LLMModelAdapter from dbgpt.model.adapter.template import ConversationAdapter, PromptType from dbgpt.model.base import ModelType -from dbgpt.model.conversation import Conversation +from dbgpt.model.llm.conversation import Conversation from dbgpt.model.parameter import ( LlamaCppModelParameters, ModelParameters, diff --git a/dbgpt/model/base.py b/dbgpt/model/base.py index 1bd5a3734..2544b70f2 100644 --- a/dbgpt/model/base.py +++ b/dbgpt/model/base.py @@ -1,13 +1,11 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- -import time from dataclasses import asdict, dataclass from datetime import datetime from enum import Enum -from typing import Any, Dict, List, Optional, TypedDict +from typing import Dict, List, Optional -from dbgpt.util.model_utils import GPUInfo from dbgpt.util.parameter_utils import ParameterDescription diff --git a/dbgpt/model/cluster/worker/default_worker.py b/dbgpt/model/cluster/worker/default_worker.py index 9023f3ced..d1f2aea5b 100644 --- a/dbgpt/model/cluster/worker/default_worker.py +++ b/dbgpt/model/cluster/worker/default_worker.py @@ -12,9 +12,9 @@ ModelOutput, ) from dbgpt.model.adapter.base import LLMModelAdapter +from dbgpt.model.adapter.loader import ModelLoader, _get_model_real_path from dbgpt.model.adapter.model_adapter import get_llm_model_adapter from dbgpt.model.cluster.worker_base import ModelWorker -from dbgpt.model.loader import ModelLoader, _get_model_real_path from dbgpt.model.parameter import ModelParameters from dbgpt.util.model_utils import _clear_model_cache, _get_current_cuda_memory from dbgpt.util.parameter_utils import EnvArgumentParser, _get_dict_from_obj diff --git a/dbgpt/model/cluster/worker/embedding_worker.py b/dbgpt/model/cluster/worker/embedding_worker.py index 06b6ce3bb..eb8cacd90 100644 --- a/dbgpt/model/cluster/worker/embedding_worker.py +++ b/dbgpt/model/cluster/worker/embedding_worker.py @@ -3,9 +3,9 @@ from dbgpt.configs.model_config import get_device from dbgpt.core import ModelMetadata +from dbgpt.model.adapter.loader import _get_model_real_path from dbgpt.model.cluster.embedding.loader import EmbeddingLoader from dbgpt.model.cluster.worker_base import ModelWorker -from dbgpt.model.loader import _get_model_real_path from dbgpt.model.parameter import ( EMBEDDING_NAME_TO_PARAMETER_CLASS_CONFIG, BaseEmbeddingModelParameters, diff --git a/dbgpt/model/cluster/worker/manager.py b/dbgpt/model/cluster/worker/manager.py index f775c9118..605f28cdd 100644 --- a/dbgpt/model/cluster/worker/manager.py +++ b/dbgpt/model/cluster/worker/manager.py @@ -8,7 +8,7 @@ import time from concurrent.futures import ThreadPoolExecutor from dataclasses import asdict -from typing import Awaitable, Callable, Dict, Iterator, List +from typing import Awaitable, Callable, Iterator from fastapi import APIRouter, FastAPI from fastapi.responses import StreamingResponse @@ -16,12 +16,7 @@ from dbgpt.component import SystemApp from dbgpt.configs.model_config import LOGDIR from dbgpt.core import ModelMetadata, ModelOutput -from dbgpt.model.base import ( - ModelInstance, - WorkerApplyOutput, - WorkerApplyType, - WorkerSupportedModel, -) +from dbgpt.model.base import ModelInstance, WorkerApplyOutput, WorkerSupportedModel from dbgpt.model.cluster.base import * from dbgpt.model.cluster.manager_base import ( WorkerManager, @@ -30,8 +25,8 @@ ) from dbgpt.model.cluster.registry import ModelRegistry from dbgpt.model.cluster.worker_base import ModelWorker -from dbgpt.model.llm_utils import list_supported_models -from dbgpt.model.parameter import ModelParameters, ModelWorkerParameters, WorkerType +from dbgpt.model.parameter import ModelWorkerParameters, WorkerType +from dbgpt.model.utils.llm_utils import list_supported_models from dbgpt.util.parameter_utils import ( EnvArgumentParser, ParameterDescription, diff --git a/dbgpt/model/compression.py b/dbgpt/model/llm/compression.py similarity index 100% rename from dbgpt/model/compression.py rename to dbgpt/model/llm/compression.py diff --git a/dbgpt/model/conversation.py b/dbgpt/model/llm/conversation.py similarity index 100% rename from dbgpt/model/conversation.py rename to dbgpt/model/llm/conversation.py diff --git a/dbgpt/model/inference.py b/dbgpt/model/llm/inference.py similarity index 99% rename from dbgpt/model/inference.py rename to dbgpt/model/llm/inference.py index fcb0d08f3..fa1525fc3 100644 --- a/dbgpt/model/inference.py +++ b/dbgpt/model/llm/inference.py @@ -18,7 +18,7 @@ TopPLogitsWarper, ) -from dbgpt.model.llm_utils import is_partial_stop, is_sentence_complete +from dbgpt.model.utils.llm_utils import is_partial_stop, is_sentence_complete def prepare_logits_processor( diff --git a/dbgpt/model/operator/__init__.py b/dbgpt/model/operator/__init__.py index b0b81d551..ec628d6b3 100644 --- a/dbgpt/model/operator/__init__.py +++ b/dbgpt/model/operator/__init__.py @@ -1,9 +1,9 @@ -from dbgpt.model.operator.llm_operator import ( +from dbgpt.model.operator.llm_operator import ( # noqa: F401 LLMOperator, MixinLLMOperator, StreamingLLMOperator, ) -from dbgpt.model.utils.chatgpt_utils import OpenAIStreamingOutputOperator +from dbgpt.model.utils.chatgpt_utils import OpenAIStreamingOutputOperator # noqa: F401 __ALL__ = [ "MixinLLMOperator", diff --git a/dbgpt/model/operator/llm_operator.py b/dbgpt/model/operator/llm_operator.py index b35d042df..00c7251c3 100644 --- a/dbgpt/model/operator/llm_operator.py +++ b/dbgpt/model/operator/llm_operator.py @@ -6,7 +6,6 @@ from dbgpt.core import LLMClient from dbgpt.core.awel import BaseOperator from dbgpt.core.operator import BaseLLM, BaseLLMOperator, BaseStreamingLLMOperator -from dbgpt.model.cluster import WorkerManagerFactory logger = logging.getLogger(__name__) @@ -19,31 +18,30 @@ class MixinLLMOperator(BaseLLM, BaseOperator, ABC): def __init__(self, default_client: Optional[LLMClient] = None, **kwargs): super().__init__(default_client) - self._default_llm_client = default_client @property def llm_client(self) -> LLMClient: if not self._llm_client: - worker_manager_factory: WorkerManagerFactory = ( - self.system_app.get_component( - ComponentType.WORKER_MANAGER_FACTORY, - WorkerManagerFactory, - default_component=None, - ) - ) - if worker_manager_factory: + try: + from dbgpt.model.cluster import WorkerManagerFactory from dbgpt.model.cluster.client import DefaultLLMClient - self._llm_client = DefaultLLMClient(worker_manager_factory.create()) - else: - if self._default_llm_client is None: - from dbgpt.model.proxy.llms.chatgpt import OpenAILLMClient - - self._default_llm_client = OpenAILLMClient() - logger.info( - f"Can't find worker manager factory, use default llm client {self._default_llm_client}." + worker_manager_factory: WorkerManagerFactory = ( + self.system_app.get_component( + ComponentType.WORKER_MANAGER_FACTORY, + WorkerManagerFactory, + default_component=None, + ) ) - self._llm_client = self._default_llm_client + if worker_manager_factory: + self._llm_client = DefaultLLMClient(worker_manager_factory.create()) + except Exception as e: + logger.warning(f"Load worker manager failed: {e}.") + if not self._llm_client: + from dbgpt.model.proxy.llms.chatgpt import OpenAILLMClient + + logger.info("Can't find worker manager factory, use OpenAILLMClient.") + self._llm_client = OpenAILLMClient() return self._llm_client diff --git a/dbgpt/model/parameter.py b/dbgpt/model/parameter.py index 6cd666a0b..810355f0d 100644 --- a/dbgpt/model/parameter.py +++ b/dbgpt/model/parameter.py @@ -6,11 +6,8 @@ from enum import Enum from typing import Dict, Optional, Tuple, Union -from dbgpt.model.conversation import conv_templates from dbgpt.util.parameter_utils import BaseParameters -suported_prompt_templates = ",".join(conv_templates.keys()) - class WorkerType(str, Enum): LLM = "llm" @@ -299,7 +296,8 @@ class ModelParameters(BaseModelParameters): prompt_template: Optional[str] = field( default=None, metadata={ - "help": f"Prompt template. If None, the prompt template is automatically determined from model path, supported template: {suported_prompt_templates}" + "help": f"Prompt template. If None, the prompt template is automatically " + f"determined from model path" }, ) max_context_size: Optional[int] = field( @@ -450,7 +448,8 @@ class ProxyModelParameters(BaseModelParameters): proxyllm_backend: Optional[str] = field( default=None, metadata={ - "help": "The model name actually pass to current proxy server url, such as gpt-3.5-turbo, gpt-4, chatglm_pro, chatglm_std and so on" + "help": "The model name actually pass to current proxy server url, such " + "as gpt-3.5-turbo, gpt-4, chatglm_pro, chatglm_std and so on" }, ) model_type: Optional[str] = field( @@ -463,13 +462,15 @@ class ProxyModelParameters(BaseModelParameters): device: Optional[str] = field( default=None, metadata={ - "help": "Device to run model. If None, the device is automatically determined" + "help": "Device to run model. If None, the device is automatically " + "determined" }, ) prompt_template: Optional[str] = field( default=None, metadata={ - "help": f"Prompt template. If None, the prompt template is automatically determined from model path, supported template: {suported_prompt_templates}" + "help": f"Prompt template. If None, the prompt template is automatically " + f"determined from model path" }, ) max_context_size: Optional[int] = field( @@ -478,7 +479,8 @@ class ProxyModelParameters(BaseModelParameters): llm_client_class: Optional[str] = field( default=None, metadata={ - "help": "The class name of llm client, such as dbgpt.model.proxy.llms.proxy_model.ProxyModel" + "help": "The class name of llm client, such as " + "dbgpt.model.proxy.llms.proxy_model.ProxyModel" }, ) diff --git a/dbgpt/model/llm_utils.py b/dbgpt/model/utils/llm_utils.py similarity index 97% rename from dbgpt/model/llm_utils.py rename to dbgpt/model/utils/llm_utils.py index fff11ba77..cef171b1e 100644 --- a/dbgpt/model/llm_utils.py +++ b/dbgpt/model/utils/llm_utils.py @@ -37,8 +37,8 @@ def list_supported_models(): def _list_supported_models( worker_type: str, model_config: Dict[str, str] ) -> List[SupportedModel]: + from dbgpt.model.adapter.loader import _get_model_real_path from dbgpt.model.adapter.model_adapter import get_llm_model_adapter - from dbgpt.model.loader import _get_model_real_path ret = [] for model_name, model_path in model_config.items(): diff --git a/dbgpt/serve/agent/team/layout/team_awel_layout.py b/dbgpt/serve/agent/team/layout/team_awel_layout.py index c22578d30..2515bad2f 100644 --- a/dbgpt/serve/agent/team/layout/team_awel_layout.py +++ b/dbgpt/serve/agent/team/layout/team_awel_layout.py @@ -67,7 +67,7 @@ async def a_run_chat( message=start_message, sender=self, reviewer=reviewer ) final_generate_context: AgentGenerateContext = await last_node.call( - call_data={"data": start_message_context} + call_data=start_message_context ) last_message = final_generate_context.rely_messages[-1] diff --git a/examples/agents/auto_plan_agent_dialogue_example.py b/examples/agents/auto_plan_agent_dialogue_example.py index f1dabde2c..7cb8d319c 100644 --- a/examples/agents/auto_plan_agent_dialogue_example.py +++ b/examples/agents/auto_plan_agent_dialogue_example.py @@ -28,7 +28,7 @@ from dbgpt.serve.agent.team.plan.team_auto_plan import AutoPlanChatManager if __name__ == "__main__": - from dbgpt.model import OpenAILLMClient + from dbgpt.model.proxy import OpenAILLMClient llm_client = OpenAILLMClient() context: AgentContext = AgentContext(conv_id="test456", llm_provider=llm_client) diff --git a/examples/agents/awel_layout_agents_chat_examples.py b/examples/agents/awel_layout_agents_chat_examples.py index 64561ed15..363f55041 100644 --- a/examples/agents/awel_layout_agents_chat_examples.py +++ b/examples/agents/awel_layout_agents_chat_examples.py @@ -30,7 +30,7 @@ test_plugin_dir = os.path.join(parent_dir, "test_files") if __name__ == "__main__": - from dbgpt.model import OpenAILLMClient + from dbgpt.model.proxy import OpenAILLMClient llm_client = OpenAILLMClient() context: AgentContext = AgentContext(conv_id="test456", llm_provider=llm_client) diff --git a/examples/agents/plugin_agent_dialogue_example.py b/examples/agents/plugin_agent_dialogue_example.py index c5a104c26..58c1f4a6d 100644 --- a/examples/agents/plugin_agent_dialogue_example.py +++ b/examples/agents/plugin_agent_dialogue_example.py @@ -24,7 +24,7 @@ from dbgpt.core.interface.llm import ModelMetadata if __name__ == "__main__": - from dbgpt.model import OpenAILLMClient + from dbgpt.model.proxy import OpenAILLMClient llm_client = OpenAILLMClient() context: AgentContext = AgentContext(conv_id="test456", llm_provider=llm_client) diff --git a/examples/agents/retrieve_summary_agent_dialogue_example.py b/examples/agents/retrieve_summary_agent_dialogue_example.py index 207f3dcaf..f5be9e92b 100644 --- a/examples/agents/retrieve_summary_agent_dialogue_example.py +++ b/examples/agents/retrieve_summary_agent_dialogue_example.py @@ -27,7 +27,7 @@ def summary_example_with_success(): - from dbgpt.model import OpenAILLMClient + from dbgpt.model.proxy import OpenAILLMClient llm_client = OpenAILLMClient() context: AgentContext = AgentContext( diff --git a/examples/agents/single_agent_dialogue_example.py b/examples/agents/single_agent_dialogue_example.py index 3b4d68ebe..9926f06e5 100644 --- a/examples/agents/single_agent_dialogue_example.py +++ b/examples/agents/single_agent_dialogue_example.py @@ -24,7 +24,7 @@ from dbgpt.core.interface.llm import ModelMetadata if __name__ == "__main__": - from dbgpt.model import OpenAILLMClient + from dbgpt.model.proxy import OpenAILLMClient llm_client = OpenAILLMClient() context: AgentContext = AgentContext(conv_id="test456", llm_provider=llm_client) diff --git a/examples/agents/single_summary_agent_dialogue_example.py b/examples/agents/single_summary_agent_dialogue_example.py index 5ffcee49e..7f11789a3 100644 --- a/examples/agents/single_summary_agent_dialogue_example.py +++ b/examples/agents/single_summary_agent_dialogue_example.py @@ -25,7 +25,7 @@ def summary_example_with_success(): - from dbgpt.model import OpenAILLMClient + from dbgpt.model.proxy import OpenAILLMClient llm_client = OpenAILLMClient() context: AgentContext = AgentContext(conv_id="summarize", llm_provider=llm_client) @@ -76,7 +76,7 @@ def summary_example_with_success(): def summary_example_with_faliure(): - from dbgpt.model import OpenAILLMClient + from dbgpt.model.proxy import OpenAILLMClient llm_client = OpenAILLMClient() context: AgentContext = AgentContext(conv_id="summarize", llm_provider=llm_client) diff --git a/examples/awel/simple_rag_retriever_example.py b/examples/awel/simple_rag_retriever_example.py index 1d48d5478..23cf29469 100644 --- a/examples/awel/simple_rag_retriever_example.py +++ b/examples/awel/simple_rag_retriever_example.py @@ -6,7 +6,7 @@ from dbgpt.configs.model_config import MODEL_PATH, PILOT_PATH from dbgpt.core.awel import DAG, HttpTrigger, JoinOperator, MapOperator -from dbgpt.model import OpenAILLMClient +from dbgpt.model.proxy import OpenAILLMClient from dbgpt.rag.chunk import Chunk from dbgpt.rag.embedding.embedding_factory import DefaultEmbeddingFactory from dbgpt.rag.operator.embedding import EmbeddingRetrieverOperator diff --git a/examples/awel/simple_rag_rewrite_example.py b/examples/awel/simple_rag_rewrite_example.py index b09929bd4..81c84a8d0 100644 --- a/examples/awel/simple_rag_rewrite_example.py +++ b/examples/awel/simple_rag_rewrite_example.py @@ -32,7 +32,7 @@ from dbgpt._private.pydantic import BaseModel, Field from dbgpt.core.awel import DAG, HttpTrigger, MapOperator -from dbgpt.model import OpenAILLMClient +from dbgpt.model.proxy import OpenAILLMClient from dbgpt.rag.operator.rewrite import QueryRewriteOperator diff --git a/examples/awel/simple_rag_summary_example.py b/examples/awel/simple_rag_summary_example.py index ca1c145fe..3b05996a7 100644 --- a/examples/awel/simple_rag_summary_example.py +++ b/examples/awel/simple_rag_summary_example.py @@ -30,7 +30,7 @@ from dbgpt._private.pydantic import BaseModel, Field from dbgpt.core.awel import DAG, HttpTrigger, MapOperator -from dbgpt.model import OpenAILLMClient +from dbgpt.model.proxy import OpenAILLMClient from dbgpt.rag.knowledge.base import KnowledgeType from dbgpt.rag.operator.knowledge import KnowledgeOperator from dbgpt.rag.operator.summary import SummaryAssemblerOperator diff --git a/examples/rag/rewrite_rag_example.py b/examples/rag/rewrite_rag_example.py index f96e53aed..a74f60b87 100644 --- a/examples/rag/rewrite_rag_example.py +++ b/examples/rag/rewrite_rag_example.py @@ -1,7 +1,7 @@ import asyncio import os -from dbgpt.model import OpenAILLMClient +from dbgpt.model.proxy import OpenAILLMClient from dbgpt.rag.retriever.rewrite import QueryRewrite """Query rewrite example. diff --git a/examples/rag/summary_extractor_example.py b/examples/rag/summary_extractor_example.py index 825c2ce76..33b4cc585 100644 --- a/examples/rag/summary_extractor_example.py +++ b/examples/rag/summary_extractor_example.py @@ -1,6 +1,6 @@ import asyncio -from dbgpt.model import OpenAILLMClient +from dbgpt.model.proxy import OpenAILLMClient from dbgpt.rag.chunk_manager import ChunkParameters from dbgpt.rag.knowledge.factory import KnowledgeFactory from dbgpt.serve.rag.assembler.summary import SummaryAssembler diff --git a/examples/sdk/simple_sdk_llm_example.py b/examples/sdk/simple_sdk_llm_example.py index ef3c43966..f141c768b 100644 --- a/examples/sdk/simple_sdk_llm_example.py +++ b/examples/sdk/simple_sdk_llm_example.py @@ -7,7 +7,7 @@ PromptBuilderOperator, RequestBuilderOperator, ) -from dbgpt.model import OpenAILLMClient +from dbgpt.model.proxy import OpenAILLMClient with DAG("simple_sdk_llm_example_dag") as dag: prompt_task = PromptBuilderOperator( @@ -20,8 +20,6 @@ if __name__ == "__main__": output = asyncio.run( - out_parse_task.call( - call_data={"data": {"dialect": "mysql", "table_name": "user"}} - ) + out_parse_task.call(call_data={"dialect": "mysql", "table_name": "user"}) ) print(f"output: \n\n{output}") diff --git a/examples/sdk/simple_sdk_llm_sql_example.py b/examples/sdk/simple_sdk_llm_sql_example.py index e42705e42..2eceefb8b 100644 --- a/examples/sdk/simple_sdk_llm_sql_example.py +++ b/examples/sdk/simple_sdk_llm_sql_example.py @@ -17,7 +17,7 @@ ) from dbgpt.datasource.operator.datasource_operator import DatasourceOperator from dbgpt.datasource.rdbms.conn_sqlite import SQLiteTempConnect -from dbgpt.model import OpenAILLMClient +from dbgpt.model.proxy import OpenAILLMClient from dbgpt.rag.operator.datasource import DatasourceRetrieverOperator @@ -144,12 +144,10 @@ def _combine_result(self, sql_result_df, model_result: Dict) -> Dict: if __name__ == "__main__": input_data = { - "data": { - "db_name": "test_db", - "dialect": "sqlite", - "top_k": 5, - "user_input": "What is the name and age of the user with age less than 18", - } + "db_name": "test_db", + "dialect": "sqlite", + "top_k": 5, + "user_input": "What is the name and age of the user with age less than 18", } output = asyncio.run(sql_result_task.call(call_data=input_data)) print(f"\nthoughts: {output.get('thoughts')}\n") diff --git a/setup.py b/setup.py index 2f318c9a0..f2fa86f33 100644 --- a/setup.py +++ b/setup.py @@ -14,6 +14,11 @@ with open("README.md", mode="r", encoding="utf-8") as fh: long_description = fh.read() +IS_DEV_MODE = os.getenv("IS_DEV_MODE", "true").lower() == "true" +# If you modify the version, please modify the version in the following files: +# dbgpt/_version.py +DB_GPT_VERSION = os.getenv("DB_GPT_VERSION", "0.4.7") + BUILD_NO_CACHE = os.getenv("BUILD_NO_CACHE", "true").lower() == "true" LLAMA_CPP_GPU_ACCELERATION = ( os.getenv("LLAMA_CPP_GPU_ACCELERATION", "true").lower() == "true" @@ -352,31 +357,41 @@ def llama_cpp_python_cuda_requires(): def core_requires(): """ - pip install db-gpt or pip install "db-gpt[core]" + pip install dbgpt or pip install "dbgpt[core]" """ setup_spec.extras["core"] = [ "aiohttp==3.8.4", "chardet==5.1.0", "importlib-resources==5.12.0", - "psutil==5.9.4", "python-dotenv==1.0.0", - "colorama==0.4.6", - "prettytable", "cachetools", + "pydantic<2,>=1", ] - # Just use by DB-GPT internal, we should find the smallest dependency set for run we core unit test. + # Simple command line dependencies + setup_spec.extras["cli"] = setup_spec.extras["core"] + [ + "prettytable", + "click", + "psutil==5.9.4", + "colorama==0.4.6", + ] + # Just use by DB-GPT internal, we should find the smallest dependency set for run + # we core unit test. # The dependency "framework" is too large for now. - setup_spec.extras["simple_framework"] = setup_spec.extras["core"] + [ + setup_spec.extras["simple_framework"] = setup_spec.extras["cli"] + [ "pydantic<2,>=1", "httpx", "jinja2", "fastapi==0.98.0", + "uvicorn", "shortuuid", - # change from fixed version 2.0.22 to variable version, because other dependencies are >=1.4, such as pydoris is <2 + # change from fixed version 2.0.22 to variable version, because other + # dependencies are >=1.4, such as pydoris is <2 "SQLAlchemy>=1.4,<3", # for cache "msgpack", - # for cache, TODO pympler has not been updated for a long time and needs to find a new toolkit. + # for cache + # TODO: pympler has not been updated for a long time and needs to + # find a new toolkit. "pympler", "sqlparse==0.4.4", "duckdb==0.8.1", @@ -418,7 +433,7 @@ def core_requires(): def knowledge_requires(): """ - pip install "db-gpt[knowledge]" + pip install "dbgpt[knowledge]" """ setup_spec.extras["knowledge"] = [ "spacy==3.5.3", @@ -435,7 +450,7 @@ def knowledge_requires(): def llama_cpp_requires(): """ - pip install "db-gpt[llama_cpp]" + pip install "dbgpt[llama_cpp]" """ setup_spec.extras["llama_cpp"] = ["llama-cpp-python"] llama_cpp_python_cuda_requires() @@ -523,7 +538,7 @@ def quantization_requires(): def all_vector_store_requires(): """ - pip install "db-gpt[vstore]" + pip install "dbgpt[vstore]" """ setup_spec.extras["vstore"] = [ "grpcio==1.47.5", # maybe delete it @@ -534,7 +549,7 @@ def all_vector_store_requires(): def all_datasource_requires(): """ - pip install "db-gpt[datasource]" + pip install "dbgpt[datasource]" """ setup_spec.extras["datasource"] = [ @@ -552,7 +567,7 @@ def all_datasource_requires(): def openai_requires(): """ - pip install "db-gpt[openai]" + pip install "dbgpt[openai]" """ setup_spec.extras["openai"] = ["tiktoken"] if BUILD_VERSION_OPENAI: @@ -567,28 +582,28 @@ def openai_requires(): def gpt4all_requires(): """ - pip install "db-gpt[gpt4all]" + pip install "dbgpt[gpt4all]" """ setup_spec.extras["gpt4all"] = ["gpt4all"] def vllm_requires(): """ - pip install "db-gpt[vllm]" + pip install "dbgpt[vllm]" """ setup_spec.extras["vllm"] = ["vllm"] def cache_requires(): """ - pip install "db-gpt[cache]" + pip install "dbgpt[cache]" """ setup_spec.extras["cache"] = ["rocksdict"] def default_requires(): """ - pip install "db-gpt[default]" + pip install "dbgpt[default]" """ setup_spec.extras["default"] = [ # "tokenizers==0.13.3", @@ -637,14 +652,46 @@ def init_install_requires(): all_requires() init_install_requires() +# Packages to exclude when IS_DEV_MODE is False +excluded_packages = ["tests", "*.tests", "*.tests.*", "examples"] + +if IS_DEV_MODE: + packages = find_packages(exclude=excluded_packages) +else: + packages = find_packages( + exclude=excluded_packages, + include=[ + "dbgpt", + "dbgpt._private", + "dbgpt._private.*", + "dbgpt.cli", + "dbgpt.cli.*", + "dbgpt.configs", + "dbgpt.configs.*", + "dbgpt.core", + "dbgpt.core.*", + "dbgpt.util", + "dbgpt.util.*", + "dbgpt.model", + "dbgpt.model.proxy", + "dbgpt.model.proxy.*", + "dbgpt.model.operator", + "dbgpt.model.operator.*", + "dbgpt.model.utils", + "dbgpt.model.utils.*", + ], + ) + setuptools.setup( - name="db-gpt", - packages=find_packages(exclude=("tests", "*.tests", "*.tests.*", "examples")), - version="0.4.5", + name="dbgpt", + packages=packages, + version=DB_GPT_VERSION, author="csunny", author_email="cfqcsunny@gmail.com", - description="DB-GPT is an experimental open-source project that uses localized GPT large models to interact with your data and environment." - " With this solution, you can be assured that there is no risk of data leakage, and your data is 100% private and secure.", + description="DB-GPT is an experimental open-source project that uses localized GPT " + "large models to interact with your data and environment." + " With this solution, you can be assured that there is no risk of data leakage, " + "and your data is 100% private and secure.", long_description=long_description, long_description_content_type="text/markdown", install_requires=setup_spec.install_requires,