github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-13 21:21:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(awel): New MessageConverter and more AWEL operators (#1039)

Browse Source
This commit is contained in:
Fangyin Cheng
2024-01-08 09:40:05 +08:00
committed by GitHub
parent 765fb181f6
commit e8861bd8fa
48 changed files with 2333 additions and 719 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
dbgpt/model/cluster/client.py
View File

@@ -1,20 +1,35 @@
from typing import AsyncIterator, List
import asyncio
from dbgpt.core.interface.llm import LLMClient, ModelRequest, ModelOutput, ModelMetadata
from dbgpt.model.parameter import WorkerType
from typing import AsyncIterator, List, Optional
from dbgpt.core.interface.llm import (
LLMClient,
MessageConverter,
ModelMetadata,
ModelOutput,
ModelRequest,
)
from dbgpt.model.cluster.manager_base import WorkerManager
from dbgpt.model.parameter import WorkerType
class DefaultLLMClient(LLMClient):
def __init__(self, worker_manager: WorkerManager):
self._worker_manager = worker_manager
async def generate(self, request: ModelRequest) -> ModelOutput:
async def generate(
self,
request: ModelRequest,
message_converter: Optional[MessageConverter] = None,
) -> ModelOutput:
request = await self.covert_message(request, message_converter)
return await self._worker_manager.generate(request.to_dict())
async def generate_stream(
self, request: ModelRequest
self,
request: ModelRequest,
message_converter: Optional[MessageConverter] = None,
) -> AsyncIterator[ModelOutput]:
request = await self.covert_message(request, message_converter)
async for output in self._worker_manager.generate_stream(request.to_dict()):
yield output