feat: filter empty content messages in llm node (#3547)

This commit is contained in:
takatost 2024-04-17 13:30:33 +08:00 committed by GitHub
commit b890c11c14
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 40 additions and 1 deletions

View file

@ -88,6 +88,14 @@ class PromptMessage(ABC, BaseModel):
content: Optional[str | list[PromptMessageContent]] = None
name: Optional[str] = None
def is_empty(self) -> bool:
"""
Check if prompt message is empty.
:return: True if prompt message is empty, False otherwise
"""
return not self.content
class UserPromptMessage(PromptMessage):
"""
@ -118,6 +126,16 @@ class AssistantPromptMessage(PromptMessage):
role: PromptMessageRole = PromptMessageRole.ASSISTANT
tool_calls: list[ToolCall] = []
def is_empty(self) -> bool:
"""
Check if prompt message is empty.
:return: True if prompt message is empty, False otherwise
"""
if not super().is_empty() and not self.tool_calls:
return False
return True
class SystemPromptMessage(PromptMessage):
"""
@ -132,3 +150,14 @@ class ToolPromptMessage(PromptMessage):
"""
role: PromptMessageRole = PromptMessageRole.TOOL
tool_call_id: str
def is_empty(self) -> bool:
"""
Check if prompt message is empty.
:return: True if prompt message is empty, False otherwise
"""
if not super().is_empty() and not self.tool_call_id:
return False
return True