diff --git a/dbgpt/app/openapi/api_v1/api_v1.py b/dbgpt/app/openapi/api_v1/api_v1.py
index 3441ab8e0..6c3dd0d86 100644
--- a/dbgpt/app/openapi/api_v1/api_v1.py
+++ b/dbgpt/app/openapi/api_v1/api_v1.py
@@ -13,11 +13,8 @@ from dbgpt._private.config import Config
from dbgpt.app.knowledge.request.request import KnowledgeSpaceRequest
from dbgpt.app.knowledge.service import KnowledgeService
from dbgpt.app.openapi.api_view_model import (
- ChatCompletionResponseStreamChoice,
- ChatCompletionStreamResponse,
ChatSceneVo,
ConversationVo,
- DeltaMessage,
MessageVo,
Result,
)
@@ -25,6 +22,11 @@ from dbgpt.app.scene import BaseChat, ChatFactory, ChatScene
from dbgpt.component import ComponentType
from dbgpt.configs.model_config import KNOWLEDGE_UPLOAD_ROOT_PATH
from dbgpt.core.awel import CommonLLMHttpRequestBody, CommonLLMHTTPRequestContext
+from dbgpt.core.schema.api import (
+ ChatCompletionResponseStreamChoice,
+ ChatCompletionStreamResponse,
+ DeltaMessage,
+)
from dbgpt.datasource.db_conn_info import DBConfig, DbTypeInfo
from dbgpt.model.base import FlatSupportedModel
from dbgpt.model.cluster import BaseModelController, WorkerManager, WorkerManagerFactory
diff --git a/dbgpt/app/openapi/api_v2.py b/dbgpt/app/openapi/api_v2.py
index 92857eadd..4ed8335d7 100644
--- a/dbgpt/app/openapi/api_v2.py
+++ b/dbgpt/app/openapi/api_v2.py
@@ -6,15 +6,6 @@ from typing import Optional
from fastapi import APIRouter, Body, Depends, HTTPException
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
-from fastchat.protocol.api_protocol import (
- ChatCompletionResponse,
- ChatCompletionResponseChoice,
- ChatCompletionResponseStreamChoice,
- ChatCompletionStreamResponse,
- ChatMessage,
- DeltaMessage,
- UsageInfo,
-)
from starlette.responses import StreamingResponse
from dbgpt.app.openapi.api_v1.api_v1 import (
@@ -26,9 +17,18 @@ from dbgpt.app.openapi.api_v1.api_v1 import (
stream_generator,
)
from dbgpt.app.scene import BaseChat, ChatScene
-from dbgpt.client.schemas import ChatCompletionRequestBody, ChatMode
+from dbgpt.client.schema import ChatCompletionRequestBody, ChatMode
from dbgpt.component import logger
from dbgpt.core.awel import CommonLLMHttpRequestBody, CommonLLMHTTPRequestContext
+from dbgpt.core.schema.api import (
+ ChatCompletionResponse,
+ ChatCompletionResponseChoice,
+ ChatCompletionResponseStreamChoice,
+ ChatCompletionStreamResponse,
+ ChatMessage,
+ DeltaMessage,
+ UsageInfo,
+)
from dbgpt.model.cluster.apiserver.api import APISettings
from dbgpt.serve.agent.agents.controller import multi_agents
from dbgpt.serve.flow.api.endpoints import get_service
diff --git a/dbgpt/app/openapi/api_view_model.py b/dbgpt/app/openapi/api_view_model.py
index cd4ccadfd..f3655e30a 100644
--- a/dbgpt/app/openapi/api_view_model.py
+++ b/dbgpt/app/openapi/api_view_model.py
@@ -89,21 +89,3 @@ class MessageVo(BaseModel):
model_name
"""
model_name: str
-
-
-class DeltaMessage(BaseModel):
- role: Optional[str] = None
- content: Optional[str] = None
-
-
-class ChatCompletionResponseStreamChoice(BaseModel):
- index: int
- delta: DeltaMessage
- finish_reason: Optional[Literal["stop", "length"]] = None
-
-
-class ChatCompletionStreamResponse(BaseModel):
- id: str = Field(default_factory=lambda: f"chatcmpl-{str(uuid.uuid1())}")
- created: int = Field(default_factory=lambda: int(time.time()))
- model: str
- choices: List[ChatCompletionResponseStreamChoice]
diff --git a/dbgpt/client/app.py b/dbgpt/client/app.py
index f90a9906b..23f2f90ea 100644
--- a/dbgpt/client/app.py
+++ b/dbgpt/client/app.py
@@ -1,9 +1,10 @@
"""App Client API."""
from typing import List
-from dbgpt.client import Client, ClientException
-from dbgpt.client.schemas import AppModel
-from dbgpt.serve.core import Result
+from dbgpt.core.schema.api import Result
+
+from .client import Client, ClientException
+from .schema import AppModel
async def get_app(client: Client, app_id: str) -> AppModel:
diff --git a/dbgpt/client/client.py b/dbgpt/client/client.py
index 60d46ba9c..5f41529f0 100644
--- a/dbgpt/client/client.py
+++ b/dbgpt/client/client.py
@@ -5,10 +5,10 @@ from typing import Any, AsyncGenerator, List, Optional, Union
from urllib.parse import urlparse
import httpx
-from fastchat.protocol.api_protocol import ChatCompletionResponse
-from dbgpt.app.openapi.api_view_model import ChatCompletionStreamResponse
-from dbgpt.client.schemas import ChatCompletionRequestBody
+from dbgpt.core.schema.api import ChatCompletionResponse, ChatCompletionStreamResponse
+
+from .schema import ChatCompletionRequestBody
CLIENT_API_PATH = "/api"
CLIENT_SERVE_PATH = "/serve"
@@ -256,14 +256,14 @@ class Client:
)
yield chat_completion_response
except Exception as e:
- yield f"data:[SERVER_ERROR]{str(e)}\n\n"
+ raise e
else:
try:
error = await response.aread()
yield json.loads(error)
except Exception as e:
- yield f"data:[SERVER_ERROR]{str(e)}\n\n"
+ raise e
async def get(self, path: str, *args):
"""Get method.
diff --git a/dbgpt/client/flow.py b/dbgpt/client/flow.py
index eda995da6..f5206b2d1 100644
--- a/dbgpt/client/flow.py
+++ b/dbgpt/client/flow.py
@@ -1,9 +1,10 @@
"""this module contains the flow client functions."""
from typing import List
-from dbgpt.client import Client, ClientException
from dbgpt.core.awel.flow.flow_factory import FlowPanel
-from dbgpt.serve.core import Result
+from dbgpt.core.schema.api import Result
+
+from .client import Client, ClientException
async def create_flow(client: Client, flow: FlowPanel) -> FlowPanel:
diff --git a/dbgpt/client/knowledge.py b/dbgpt/client/knowledge.py
index 323b51fcb..f8b3404ec 100644
--- a/dbgpt/client/knowledge.py
+++ b/dbgpt/client/knowledge.py
@@ -2,9 +2,10 @@
import json
from typing import List
-from dbgpt.client import Client, ClientException
-from dbgpt.client.schemas import DocumentModel, SpaceModel, SyncModel
-from dbgpt.serve.core import Result
+from dbgpt.core.schema.api import Result
+
+from .client import Client, ClientException
+from .schema import DocumentModel, SpaceModel, SyncModel
async def create_space(client: Client, space_model: SpaceModel) -> SpaceModel:
diff --git a/dbgpt/client/schemas.py b/dbgpt/client/schema.py
similarity index 99%
rename from dbgpt/client/schemas.py
rename to dbgpt/client/schema.py
index aab0f0211..de33e5351 100644
--- a/dbgpt/client/schemas.py
+++ b/dbgpt/client/schema.py
@@ -4,8 +4,8 @@ from enum import Enum
from typing import List, Optional, Union
from fastapi import File, UploadFile
-from pydantic import BaseModel, Field
+from dbgpt._private.pydantic import BaseModel, Field
from dbgpt.agent.resource.resource_api import AgentResource
from dbgpt.rag.chunk_manager import ChunkParameters
diff --git a/dbgpt/core/schema/__init__.py b/dbgpt/core/schema/__init__.py
new file mode 100644
index 000000000..0acfbeaa4
--- /dev/null
+++ b/dbgpt/core/schema/__init__.py
@@ -0,0 +1 @@
+"""Module for common schemas."""
diff --git a/dbgpt/core/schema/api.py b/dbgpt/core/schema/api.py
new file mode 100644
index 000000000..647987b0f
--- /dev/null
+++ b/dbgpt/core/schema/api.py
@@ -0,0 +1,116 @@
+"""API schema module."""
+
+import time
+import uuid
+from typing import Any, Generic, List, Literal, Optional, TypeVar
+
+from dbgpt._private.pydantic import BaseModel, Field
+
+T = TypeVar("T")
+
+
+class Result(BaseModel, Generic[T]):
+ """Common result entity for API response."""
+
+ success: bool = Field(
+ ..., description="Whether it is successful, True: success, False: failure"
+ )
+ err_code: str | None = Field(None, description="Error code")
+ err_msg: str | None = Field(None, description="Error message")
+ data: T | None = Field(None, description="Return data")
+
+ @staticmethod
+ def succ(data: T) -> "Result[T]":
+ """Build a successful result entity.
+
+ Args:
+ data (T): Return data
+
+ Returns:
+ Result[T]: Result entity
+ """
+ return Result(success=True, err_code=None, err_msg=None, data=data)
+
+ @staticmethod
+ def failed(msg: str, err_code: Optional[str] = "E000X") -> "Result[Any]":
+ """Build a failed result entity.
+
+ Args:
+ msg (str): Error message
+ err_code (Optional[str], optional): Error code. Defaults to "E000X".
+ """
+ return Result(success=False, err_code=err_code, err_msg=msg, data=None)
+
+
+class DeltaMessage(BaseModel):
+ """Delta message entity for chat completion response."""
+
+ role: Optional[str] = None
+ content: Optional[str] = None
+
+
+class ChatCompletionResponseStreamChoice(BaseModel):
+ """Chat completion response choice entity."""
+
+ index: int = Field(..., description="Choice index")
+ delta: DeltaMessage = Field(..., description="Delta message")
+ finish_reason: Optional[Literal["stop", "length"]] = Field(
+ None, description="Finish reason"
+ )
+
+
+class ChatCompletionStreamResponse(BaseModel):
+ """Chat completion response stream entity."""
+
+ id: str = Field(
+ default_factory=lambda: f"chatcmpl-{str(uuid.uuid1())}", description="Stream ID"
+ )
+ created: int = Field(
+ default_factory=lambda: int(time.time()), description="Created time"
+ )
+ model: str = Field(..., description="Model name")
+ choices: List[ChatCompletionResponseStreamChoice] = Field(
+ ..., description="Chat completion response choices"
+ )
+
+
+class ChatMessage(BaseModel):
+ """Chat message entity."""
+
+ role: str = Field(..., description="Role of the message")
+ content: str = Field(..., description="Content of the message")
+
+
+class UsageInfo(BaseModel):
+ """Usage info entity."""
+
+ prompt_tokens: int = Field(0, description="Prompt tokens")
+ total_tokens: int = Field(0, description="Total tokens")
+ completion_tokens: Optional[int] = Field(0, description="Completion tokens")
+
+
+class ChatCompletionResponseChoice(BaseModel):
+ """Chat completion response choice entity."""
+
+ index: int = Field(..., description="Choice index")
+ message: ChatMessage = Field(..., description="Chat message")
+ finish_reason: Optional[Literal["stop", "length"]] = Field(
+ None, description="Finish reason"
+ )
+
+
+class ChatCompletionResponse(BaseModel):
+ """Chat completion response entity."""
+
+ id: str = Field(
+ default_factory=lambda: f"chatcmpl-{str(uuid.uuid1())}", description="Stream ID"
+ )
+ object: str = "chat.completion"
+ created: int = Field(
+ default_factory=lambda: int(time.time()), description="Created time"
+ )
+ model: str = Field(..., description="Model name")
+ choices: List[ChatCompletionResponseChoice] = Field(
+ ..., description="Chat completion response choices"
+ )
+ usage: UsageInfo = Field(..., description="Usage info")
diff --git a/dbgpt/model/utils/chatgpt_utils.py b/dbgpt/model/utils/chatgpt_utils.py
index d84861dd5..3dbb3f8a3 100644
--- a/dbgpt/model/utils/chatgpt_utils.py
+++ b/dbgpt/model/utils/chatgpt_utils.py
@@ -16,14 +16,7 @@ from typing import (
from dbgpt._private.pydantic import model_to_json
from dbgpt.core.awel import TransformStreamAbsOperator
-from dbgpt.core.awel.flow import (
- IOField,
- OperatorCategory,
- OperatorType,
- Parameter,
- ResourceCategory,
- ViewMetadata,
-)
+from dbgpt.core.awel.flow import IOField, OperatorCategory, OperatorType, ViewMetadata
from dbgpt.core.interface.llm import ModelOutput
from dbgpt.core.operators import BaseLLM
@@ -217,7 +210,8 @@ async def _to_openai_stream(
import json
import shortuuid
- from fastchat.protocol.openai_api_protocol import (
+
+ from dbgpt.core.schema.api import (
ChatCompletionResponseStreamChoice,
ChatCompletionStreamResponse,
DeltaMessage,
diff --git a/dbgpt/serve/core/schemas.py b/dbgpt/serve/core/schemas.py
index 5add79f30..698ac6805 100644
--- a/dbgpt/serve/core/schemas.py
+++ b/dbgpt/serve/core/schemas.py
@@ -1,12 +1,12 @@
import logging
import sys
-from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar
+from typing import TYPE_CHECKING
from fastapi import HTTPException, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
-from dbgpt._private.pydantic import BaseModel, Field
+from dbgpt.core.schema.api import Result
if sys.version_info < (3, 11):
try:
@@ -18,40 +18,6 @@ if TYPE_CHECKING:
from fastapi import FastAPI
logger = logging.getLogger(__name__)
-T = TypeVar("T")
-
-
-class Result(BaseModel, Generic[T]):
- """Common result entity class"""
-
- success: bool = Field(
- ..., description="Whether it is successful, True: success, False: failure"
- )
- err_code: str | None = Field(None, description="Error code")
- err_msg: str | None = Field(None, description="Error message")
- data: T | None = Field(None, description="Return data")
-
- @staticmethod
- def succ(data: T) -> "Result[T]":
- """Build a successful result entity
-
- Args:
- data (T): Return data
-
- Returns:
- Result[T]: Result entity
- """
- return Result(success=True, err_code=None, err_msg=None, data=data)
-
- @staticmethod
- def failed(msg: str, err_code: Optional[str] = "E000X") -> "Result[Any]":
- """Build a failed result entity
-
- Args:
- msg (str): Error message
- err_code (Optional[str], optional): Error code. Defaults to "E000X".
- """
- return Result(success=False, err_code=err_code, err_msg=msg, data=None)
async def validation_exception_handler(
diff --git a/dbgpt/serve/flow/service/service.py b/dbgpt/serve/flow/service/service.py
index 6deeb8985..a8811fa81 100644
--- a/dbgpt/serve/flow/service/service.py
+++ b/dbgpt/serve/flow/service/service.py
@@ -485,9 +485,7 @@ async def _chat_with_dag_task(
if OpenAIStreamingOutputOperator and isinstance(
task, OpenAIStreamingOutputOperator
):
- from fastchat.protocol.openai_api_protocol import (
- ChatCompletionResponseStreamChoice,
- )
+ from dbgpt.core.schema.api import ChatCompletionResponseStreamChoice
previous_text = ""
async for output in await task.call_stream(request):
diff --git a/docs/docs/api/app.md b/docs/docs/api/app.md
index 013b4e297..f2cbec390 100644
--- a/docs/docs/api/app.md
+++ b/docs/docs/api/app.md
@@ -50,11 +50,12 @@ APP_ID="{YOUR_APP_ID}"
client = Client(api_key=DBGPT_API_KEY)
async for data in client.chat_stream(
- messages="Introduce AWEL",
- model="chatgpt_proxyllm",
- chat_mode="chat_app",
- chat_param=APP_ID):
- print(data)
+ messages="Introduce AWEL",
+ model="chatgpt_proxyllm",
+ chat_mode="chat_app",
+ chat_param=APP_ID
+):
+ print(data)
```
diff --git a/docs/docs/api/chat.md b/docs/docs/api/chat.md
index 6cf3d2621..da64c99a8 100644
--- a/docs/docs/api/chat.md
+++ b/docs/docs/api/chat.md
@@ -45,6 +45,7 @@ import TabItem from '@theme/TabItem';
from dbgpt.client import Client
DBGPT_API_KEY = "dbgpt"
+client = Client(api_key=DBGPT_API_KEY)
async for data in client.chat_stream(
model="chatgpt_proxyllm",
diff --git a/docs/docs/api/flow.md b/docs/docs/api/flow.md
index 7e63e93cf..88be326d2 100644
--- a/docs/docs/api/flow.md
+++ b/docs/docs/api/flow.md
@@ -49,11 +49,12 @@ FLOW_ID="{YOUR_FLOW_ID}"
client = Client(api_key=DBGPT_API_KEY)
async for data in client.chat_stream(
- messages="Introduce AWEL",
- model="chatgpt_proxyllm",
- chat_mode="chat_flow",
- chat_param=FLOW_ID):
- print(data)
+ messages="Introduce AWEL",
+ model="chatgpt_proxyllm",
+ chat_mode="chat_flow",
+ chat_param=FLOW_ID
+):
+ print(data)
```
diff --git a/docs/docs/api/knowledge.md b/docs/docs/api/knowledge.md
index b19be1478..8810f05fb 100644
--- a/docs/docs/api/knowledge.md
+++ b/docs/docs/api/knowledge.md
@@ -49,11 +49,12 @@ SPACE_NAME="{YOUR_SPACE_NAME}"
client = Client(api_key=DBGPT_API_KEY)
async for data in client.chat_stream(
- messages="Introduce AWEL",
- model="chatgpt_proxyllm",
- chat_mode="chat_knowledge",
- chat_param=SPACE_NAME):
- print(data)
+ messages="Introduce AWEL",
+ model="chatgpt_proxyllm",
+ chat_mode="chat_knowledge",
+ chat_param=SPACE_NAME
+):
+ print(data)
```
@@ -343,21 +344,20 @@ POST /api/v2/serve/knowledge/spaces
-
```python
from dbgpt.client import Client
from dbgpt.client.knowledge import create_space
-from dbgpt.client.schemas import SpaceModel
+from dbgpt.client.schema import SpaceModel
DBGPT_API_KEY = "dbgpt"
client = Client(api_key=DBGPT_API_KEY)
-res = await create_space(client,SpaceModel(
- name="test_space",
- vector_type="Chroma",
- desc="for client space",
- owner="dbgpt"))
-
+res = await create_space(client, SpaceModel(
+ name="test_space",
+ vector_type="Chroma",
+ desc="for client space",
+ owner="dbgpt"
+))
```
@@ -420,20 +420,20 @@ PUT /api/v2/serve/knowledge/spaces
-
```python
from dbgpt.client import Client
from dbgpt.client.knowledge import update_space
-from dbgpt.client.schemas import SpaceModel
+from dbgpt.client.schema import SpaceModel
DBGPT_API_KEY = "dbgpt"
client = Client(api_key=DBGPT_API_KEY)
res = await update_space(client, SpaceModel(
- name="test_space",
- vector_type="Chroma",
- desc="for client space update",
- owner="dbgpt"))
+ name="test_space",
+ vector_type="Chroma",
+ desc="for client space update",
+ owner="dbgpt"
+))
```
diff --git a/examples/client/knowledge_crud_example.py b/examples/client/knowledge_crud_example.py
index f3533bea2..879984255 100644
--- a/examples/client/knowledge_crud_example.py
+++ b/examples/client/knowledge_crud_example.py
@@ -66,7 +66,7 @@ import asyncio
from dbgpt.client import Client
from dbgpt.client.knowledge import create_space
-from dbgpt.client.schemas import SpaceModel
+from dbgpt.client.schema import SpaceModel
async def main():