processors(rtvi): linting and make send_error() public

This commit is contained in:
Aleix Conchillo Flaqué
2024-08-13 11:21:51 -07:00
parent 6b53c6add3
commit 87525b085e

View File

@@ -177,6 +177,7 @@ class RTVIActionResponse(BaseModel):
id: str id: str
data: RTVIActionResponseData data: RTVIActionResponseData
class RTVIBotReadyData(BaseModel): class RTVIBotReadyData(BaseModel):
version: str version: str
config: List[RTVIServiceConfig] config: List[RTVIServiceConfig]
@@ -201,7 +202,7 @@ class RTVILLMFunctionCallMessage(BaseModel):
class RTVILLMFunctionCallStartMessageData(BaseModel): class RTVILLMFunctionCallStartMessageData(BaseModel):
function_name: str function_name: str
class RTVILLMFunctionCallStartMessage(BaseModel): class RTVILLMFunctionCallStartMessage(BaseModel):
@@ -209,6 +210,7 @@ class RTVILLMFunctionCallStartMessage(BaseModel):
type: Literal["llm-function-call-start"] = "llm-function-call-start" type: Literal["llm-function-call-start"] = "llm-function-call-start"
data: RTVILLMFunctionCallStartMessageData data: RTVILLMFunctionCallStartMessageData
class RTVILLMFunctionCallResultData(BaseModel): class RTVILLMFunctionCallResultData(BaseModel):
function_name: str function_name: str
tool_call_id: str tool_call_id: str
@@ -260,10 +262,35 @@ class RTVIProcessor(FrameProcessor):
def register_service(self, service: RTVIService): def register_service(self, service: RTVIService):
self._registered_services[service.name] = service self._registered_services[service.name] = service
async def interrupt_bot(self): async def interrupt_bot(self):
await self.push_frame(BotInterruptionFrame(), FrameDirection.UPSTREAM) await self.push_frame(BotInterruptionFrame(), FrameDirection.UPSTREAM)
async def send_error(self, error: str):
message = RTVIError(data=RTVIErrorData(message=error))
frame = TransportMessageFrame(message=message.model_dump(exclude_none=True))
await self.push_frame(frame)
async def handle_function_call(
self,
function_name: str,
tool_call_id: str,
arguments: dict,
context,
result_callback):
fn = RTVILLMFunctionCallMessageData(
function_name=function_name,
tool_call_id=tool_call_id,
args=arguments)
message = RTVILLMFunctionCallMessage(data=fn)
frame = TransportMessageFrame(message=message.model_dump())
await self.push_frame(frame)
async def handle_function_call_start(self, function_name: str):
fn = RTVILLMFunctionCallStartMessageData(function_name=function_name)
message = RTVILLMFunctionCallStartMessage(data=fn)
frame = TransportMessageFrame(message=message.model_dump())
await self.push_frame(frame)
async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM): async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM):
if isinstance(frame, SystemFrame): if isinstance(frame, SystemFrame):
await super().push_frame(frame, direction) await super().push_frame(frame, direction)
@@ -302,18 +329,6 @@ class RTVIProcessor(FrameProcessor):
else: else:
await self.push_frame(frame, direction) await self.push_frame(frame, direction)
async def handle_function_call(self, function_name, tool_call_id, arguments, context, result_callback):
fn = RTVILLMFunctionCallMessageData(function_name=function_name, tool_call_id=tool_call_id, args=arguments)
message = RTVILLMFunctionCallMessage(data=fn)
frame = TransportMessageFrame(message=message.model_dump())
await self.push_frame(frame)
async def handle_function_call_start(self, function_name):
fn = RTVILLMFunctionCallStartMessageData(function_name=function_name)
message = RTVILLMFunctionCallStartMessage(data=fn)
frame = TransportMessageFrame(message=message.model_dump())
await self.push_frame(frame)
async def cleanup(self): async def cleanup(self):
if self._pipeline: if self._pipeline:
await self._pipeline.cleanup() await self._pipeline.cleanup()
@@ -385,7 +400,7 @@ class RTVIProcessor(FrameProcessor):
try: try:
message = RTVIMessage.model_validate(frame.message) message = RTVIMessage.model_validate(frame.message)
except ValidationError as e: except ValidationError as e:
await self._send_error(f"Invalid incoming message: {e}") await self.send_error(f"Invalid incoming message: {e}")
logger.warning(f"Invalid incoming message: {e}") logger.warning(f"Invalid incoming message: {e}")
return return
@@ -467,10 +482,10 @@ class RTVIProcessor(FrameProcessor):
async def _handle_function_call_result(self, data): async def _handle_function_call_result(self, data):
frame = FunctionCallResultFrame( frame = FunctionCallResultFrame(
function_name=data.function_name, function_name=data.function_name,
tool_call_id=data.tool_call_id, tool_call_id=data.tool_call_id,
arguments=data.arguments, arguments=data.arguments,
result=data.result) result=data.result)
await self.push_frame(frame) await self.push_frame(frame)
async def _handle_action(self, request_id: str, data: RTVIActionRun): async def _handle_action(self, request_id: str, data: RTVIActionRun):
@@ -496,11 +511,6 @@ class RTVIProcessor(FrameProcessor):
frame = TransportMessageFrame(message=message.model_dump(exclude_none=True)) frame = TransportMessageFrame(message=message.model_dump(exclude_none=True))
await self.push_frame(frame) await self.push_frame(frame)
async def _send_error(self, error: str):
message = RTVIError(data=RTVIErrorData(message=error))
frame = TransportMessageFrame(message=message.model_dump(exclude_none=True))
await self.push_frame(frame)
async def _send_error_response(self, id: str, error: str): async def _send_error_response(self, id: str, error: str):
message = RTVIErrorResponse(id=id, data=RTVIErrorResponseData(error=error)) message = RTVIErrorResponse(id=id, data=RTVIErrorResponseData(error=error))
frame = TransportMessageFrame(message=message.model_dump(exclude_none=True)) frame = TransportMessageFrame(message=message.model_dump(exclude_none=True))