wehub-resource-sync
70 lines
2.9 KiB
Python
70 lines
2.9 KiB
Python
from sqlalchemy import Column, Integer, Text, String, asc
|
|
from sqlalchemy.orm import Session
|
|
|
|
from superagi.models.agent_execution import AgentExecution
|
|
from superagi.models.base_model import DBBaseModel
|
|
|
|
|
|
class AgentExecutionFeed(DBBaseModel):
|
|
"""
|
|
Feed of the agent execution.
|
|
|
|
Attributes:
|
|
id (int): The unique identifier of the agent execution feed.
|
|
agent_execution_id (int): The identifier of the associated agent execution.
|
|
agent_id (int): The identifier of the associated agent.
|
|
feed (str): The feed content.
|
|
role (str): The role of the feed entry. Possible values: 'system', 'user', or 'assistant'.
|
|
extra_info (str): Additional information related to the feed entry.
|
|
"""
|
|
|
|
__tablename__ = 'agent_execution_feeds'
|
|
|
|
id = Column(Integer, primary_key=True)
|
|
agent_execution_id = Column(Integer)
|
|
agent_id = Column(Integer)
|
|
feed = Column(Text)
|
|
role = Column(String)
|
|
extra_info = Column(String)
|
|
feed_group_id = Column(String)
|
|
error_message = Column(String)
|
|
|
|
def __repr__(self):
|
|
"""
|
|
Returns a string representation of the AgentExecutionFeed object.
|
|
|
|
Returns:
|
|
str: String representation of the AgentExecutionFeed.
|
|
"""
|
|
|
|
return f"AgentExecutionFeed(id={self.id}, " \
|
|
f"agent_execution_id={self.agent_execution_id}, " \
|
|
f"feed='{self.feed}', role='{self.role}', extra_info='{self.extra_info}', feed_group_id='{self.feed_group_id}')"
|
|
|
|
@classmethod
|
|
def get_last_tool_response(cls, session: Session, agent_execution_id: int, tool_name: str = None):
|
|
agent_execution_feeds = session.query(AgentExecutionFeed).filter(
|
|
AgentExecutionFeed.agent_execution_id == agent_execution_id,
|
|
AgentExecutionFeed.role == "system").order_by(AgentExecutionFeed.created_at.desc()).all()
|
|
|
|
for agent_execution_feed in agent_execution_feeds:
|
|
if tool_name and not agent_execution_feed.feed.startswith("Tool " + tool_name):
|
|
continue
|
|
if agent_execution_feed.feed.startswith("Tool"):
|
|
return agent_execution_feed.feed
|
|
return ""
|
|
|
|
@classmethod
|
|
def fetch_agent_execution_feeds(cls, session, agent_execution_id: int):
|
|
agent_execution = AgentExecution.find_by_id(session, agent_execution_id)
|
|
agent_feeds = session.query(AgentExecutionFeed.role, AgentExecutionFeed.feed, AgentExecutionFeed.id) \
|
|
.filter(AgentExecutionFeed.agent_execution_id == agent_execution_id,
|
|
AgentExecutionFeed.feed_group_id == agent_execution.current_feed_group_id) \
|
|
.order_by(asc(AgentExecutionFeed.created_at)) \
|
|
.all()
|
|
# return entire feed if it is not default feed. Default feed has prompt in the first 2 entries.
|
|
if agent_execution.current_feed_group_id != "DEFAULT":
|
|
return agent_feeds
|
|
else:
|
|
return agent_feeds[2:]
|