ReMe/reme/extension/simple_chat.py

60 lines
2.2 KiB
Python

"""Simple chat for test."""
from loguru import logger
from ..core.enumeration import Role
from ..core.op import BaseTool
from ..core.schema import Message, ToolCall
class SimpleChat(BaseTool):
"""Simple chat agent that handles non-streaming conversations."""
def _build_tool_call(self) -> ToolCall:
return ToolCall(
**{
"description": "simple chat agent",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "query",
},
"messages": {
"type": "array",
"items": {
"type": "object",
"properties": {
"role": {
"type": "string",
"description": "role",
},
"content": {
"type": "string",
"description": "content",
},
},
"required": ["role", "content"],
},
},
},
"required": [],
},
},
)
async def execute(self):
if "query" in self.context:
messages = [
Message(role=Role.SYSTEM, content="You are a helpful assistant."),
Message(role=Role.USER, content=self.context.query),
]
elif "messages" in self.context:
messages = [Message(**m) if isinstance(m, dict) else m for m in self.context.messages if m]
else:
raise ValueError("query or messages must be provided!")
logger.info(f"messages={messages}")
assistant_message = await self.llm.chat(messages=messages)
logger.info(f"assistant_message={assistant_message.simple_dump()}")
return assistant_message.content