ReMe/experiencemaker/module/agent_wrapper/simple_agent_wrapper.py
2025-06-11 23:13:46 +08:00

18 lines
766 B
Python

from experiencemaker.module.agent_wrapper.agent_wrapper_mixin import AgentWrapperMixin
from experiencemaker.module.agent_wrapper.simple_agent import SimpleAgent
from experiencemaker.schema.trajectory import Trajectory
class SimpleAgentWrapper(SimpleAgent, AgentWrapperMixin):
def execute(self, query: str, workspace_id: str = None, **kwargs) -> Trajectory:
trajectory = Trajectory(query=query)
context_msg = self.context_generator.execute(trajectory=trajectory, workspace_id=workspace_id)
new_query = f"{context_msg.content}\n\nUser Question\n{query}"
messages = self.run(new_query)
trajectory.steps = messages
trajectory.answer = messages[-1].content
trajectory.done = True
return trajectory