Entity fix

superagi/agent/output_handler.py

@@ -100,7 +100,7 @@ def handle_tool_response(self, session, assistant_reply):
         action = self.output_parser.parse(assistant_reply)
         agent = session.query(Agent).filter(Agent.id == self.agent_config["agent_id"]).first()
         organisation = agent.get_agent_organisation(session)
-        tool_executor = ToolExecutor(organisation_id=organisation.id, agent_id=agent.id, tools=self.tools)
+        tool_executor = ToolExecutor(organisation_id=organisation.id, agent_id=agent.id, tools=self.tools, agent_execution_id=self.agent_execution_id)
         return tool_executor.execute(session, action.name, action.args)
 
     def _check_permission_in_restricted_mode(self, session, assistant_reply: str):

superagi/agent/tool_executor.py

@@ -9,10 +9,11 @@ class ToolExecutor:
     """Executes the tool with the given args."""
     FINISH = "finish"
 
-    def __init__(self, organisation_id: int, agent_id: int, tools: list):
+    def __init__(self, organisation_id: int, agent_id: int, tools: list, agent_execution_id: int):
         self.organisation_id = organisation_id
         self.agent_id = agent_id
         self.tools = tools
+        self.agent_execution_id = agent_execution_id
 
     def execute(self, session, tool_name, tool_args):
         """Executes the tool with the given args.
@@ -31,7 +32,7 @@ def execute(self, session, tool_name, tool_args):
             status = "SUCCESS"
             tool = tools[tool_name]
             retry = False
-            EventHandler(session=session).create_event('tool_used', {'tool_name': tool_name}, self.agent_id,
+            EventHandler(session=session).create_event('tool_used', {'tool_name': tool_name, 'agent_execution_id': self.agent_execution_id}, self.agent_id,
                                                        self.organisation_id),
             try:
                 parsed_args = self.clean_tool_args(tool_args)

superagi/apm/knowledge_handler.py

@@ -8,6 +8,7 @@
 from sqlalchemy.orm import aliased
 from superagi.models.agent_config import AgentConfiguration
 import pytz
+from datetime import datetime
 
 
 class KnowledgeHandler:
@@ -56,78 +57,61 @@ def get_knowledge_usage_by_name(self, knowledge_name: str) -> Dict[str, Dict[str
         return knowledge_data
 
 
-    def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Union[str, int]]]:
-        
+    def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Union[str, int, List[str]]]]:
+
         is_knowledge_valid = self.session.query(Knowledges.id).filter_by(name=knowledge_name).filter(Knowledges.organisation_id == self.organisation_id).first()
 
         if not is_knowledge_valid:
             raise HTTPException(status_code=404, detail="Knowledge not found")
 
-        event_knowledge_picked = self.session.query(
-            Event.agent_id,
-            label('created_at', func.max(Event.created_at)),
-            label('event_name', func.max(Event.event_name))
-        ).filter(
+        knowledge_events = self.session.query(Event).filter(
             Event.org_id == self.organisation_id,
             Event.event_name == 'knowledge_picked',
             Event.event_property['knowledge_name'].astext == knowledge_name
-        ).group_by(Event.agent_id).subquery()
+        ).all()
 
-        event_run = self.session.query(
-            Event.agent_id,
-            label('tokens_consumed', func.sum(Event.event_property['tokens_consumed'].astext.cast(Integer))),
-            label('calls', func.sum(Event.event_property['calls'].astext.cast(Integer)))
-        ).filter(
-            Event.org_id == self.organisation_id,
-            or_(Event.event_name == 'run_completed', Event.event_name == 'run_iteration_limit_crossed')
-        ).group_by(Event.agent_id).subquery()
+        knowledge_events = [ke for ke in knowledge_events if 'agent_execution_id' in ke.event_property]
 
-        event_run_created = self.session.query(
-            Event.agent_id,
-            label('agent_execution_name', func.max(Event.event_property['agent_execution_name'].astext))
-        ).filter(
+        event_runs = self.session.query(Event).filter(
             Event.org_id == self.organisation_id,
-            Event.event_name == 'run_created'
-        ).group_by(Event.agent_id).subquery()
+            or_(Event.event_name == 'run_completed', Event.event_name == 'run_iteration_limit_crossed')
+        ).all()
 
-        event_agent_created = self.session.query(
-            Event.agent_id,
-            label('agent_name', func.max(Event.event_property['agent_name'].astext)),
-            label('model', func.max(Event.event_property['model'].astext))
-        ).filter(
+        agent_created_events = self.session.query(Event).filter(
             Event.org_id == self.organisation_id,
             Event.event_name == 'agent_created'
-        ).group_by(Event.agent_id).subquery()
-
-
-        result = self.session.query(
-            event_knowledge_picked.c.agent_id,
-            event_knowledge_picked.c.created_at,
-            event_knowledge_picked.c.event_name,
-            event_run.c.tokens_consumed,
-            event_run.c.calls,
-            event_run_created.c.agent_execution_name,
-            event_agent_created.c.agent_name,
-            event_agent_created.c.model
-        ).join(
-            event_run, event_knowledge_picked.c.agent_id == event_run.c.agent_id
-        ).join(
-            event_run_created, event_knowledge_picked.c.agent_id == event_run_created.c.agent_id
-        ).join(
-            event_agent_created, event_knowledge_picked.c.agent_id == event_agent_created.c.agent_id
         ).all()
 
-        user_timezone = AgentConfiguration.get_agent_config_by_key_and_agent_id(session= self.session,key= 'user_timezone', agent_id= Event.agent_id)
-        if user_timezone.value is None:
-            user_timezone.value = 'GMT'
-
-        return [{
-            'agent_id': row.agent_id,
-            'created_at': row.created_at.astimezone(pytz.timezone(user_timezone.value)).strftime("%d %B %Y %H:%M"),
-            'event_name': row.event_name,
-            'tokens_consumed': row.tokens_consumed,
-            'calls': row.calls,
-            'agent_execution_name': row.agent_execution_name,
-            'agent_name': row.agent_name,
-            'model': row.model
-        } for row in result]
+        results = []
+
+        for knowledge_event in knowledge_events:
+            agent_execution_id = knowledge_event.event_property['agent_execution_id']
+
+            event_run = next((er for er in event_runs if er.agent_id == knowledge_event.agent_id and er.event_property['agent_execution_id'] == agent_execution_id), None)
+            agent_created_event = next((ace for ace in agent_created_events if ace.agent_id == knowledge_event.agent_id), None)
+            try:
+                user_timezone = AgentConfiguration.get_agent_config_by_key_and_agent_id(session=self.session, key='user_timezone', agent_id=knowledge_event.agent_id)
+                if user_timezone and user_timezone.value != 'None':
+                    tz = pytz.timezone(user_timezone.value)
+                else:
+                    tz = pytz.timezone('GMT')       
+            except AttributeError:
+                tz = pytz.timezone('GMT')
+
+            if event_run and agent_created_event:
+                actual_time = knowledge_event.created_at.astimezone(tz).strftime("%d %B %Y %H:%M")
+
+                result_dict = {
+                    'agent_execution_id': agent_execution_id,
+                    'created_at': actual_time,
+                    'tokens_consumed': event_run.event_property['tokens_consumed'],
+                    'calls': event_run.event_property['calls'],
+                    'agent_execution_name': event_run.event_property['name'],
+                    'agent_name': agent_created_event.event_property['agent_name'],
+                    'model': agent_created_event.event_property['model']
+                }
+                if agent_execution_id not in [i['agent_execution_id'] for i in results]:
+                    results.append(result_dict)
+
+        results = sorted(results, key=lambda x: datetime.strptime(x['created_at'], '%d %B %Y %H:%M'), reverse=True)
+        return results

superagi/apm/tools_handler.py

@@ -1,5 +1,5 @@
 from typing import List, Dict, Union
-from sqlalchemy import func, distinct
+from sqlalchemy import func, distinct, and_
 from sqlalchemy.orm import Session
 from sqlalchemy import Integer
 from fastapi import HTTPException
@@ -8,6 +8,7 @@
 from superagi.models.toolkit import Toolkit
 from sqlalchemy import or_
 from sqlalchemy.sql import label
+from datetime import datetime
 from superagi.models.agent_config import AgentConfiguration
 import pytz
 
@@ -87,93 +88,79 @@ def get_tool_usage_by_name(self, tool_name: str) -> Dict[str, Dict[str, int]]:
             }
 
         return tool_data
+
 
-    def get_tool_events_by_name(self, tool_name: str) -> List[Dict[str, Union[str, int, List[str]]]]:    
-
+    def get_tool_events_by_name(self, tool_name: str) -> List[Dict[str, Union[str, int, List[str]]]]:
         is_tool_name_valid = self.session.query(Tool).filter_by(name=tool_name).first()
 
         if not is_tool_name_valid:
             raise HTTPException(status_code=404, detail="Tool not found")
 
         formatted_tool_name = tool_name.lower().replace(" ", "")
 
-        event_tool_used = self.session.query(
-            Event.agent_id,
-            label('created_at', func.max(Event.created_at)),
-            label('event_name', func.max(Event.event_name))
-        ).filter(
+        tool_events = self.session.query(Event).filter(
             Event.org_id == self.organisation_id,
             Event.event_name == 'tool_used',
             Event.event_property['tool_name'].astext == formatted_tool_name
-        ).group_by(Event.agent_id).subquery()
+        ).all()
 
-        event_run = self.session.query(
-            Event.agent_id,
-            label('tokens_consumed', func.sum(Event.event_property['tokens_consumed'].astext.cast(Integer))),
-            label('calls', func.sum(Event.event_property['calls'].astext.cast(Integer)))
-        ).filter(
-            Event.org_id == self.organisation_id,
-            or_(Event.event_name == 'run_completed', Event.event_name == 'run_iteration_limit_crossed')
-        ).group_by(Event.agent_id).subquery()
+        tool_events = [te for te in tool_events if 'agent_execution_id' in te.event_property]
 
-        event_run_created = self.session.query(
-            Event.agent_id,
-            label('agent_execution_name', func.max(Event.event_property['agent_execution_name'].astext))
-        ).filter(
+        event_runs = self.session.query(Event).filter(
             Event.org_id == self.organisation_id,
-            Event.event_name == 'run_created'
-        ).group_by(Event.agent_id).subquery()
+            or_(Event.event_name == 'run_completed', Event.event_name == 'run_iteration_limit_crossed')
+        ).all()
 
-        event_agent_created = self.session.query(
-            Event.agent_id,
-            label('agent_name', func.max(Event.event_property['agent_name'].astext)),
-            label('model', func.max(Event.event_property['model'].astext))
-        ).filter(
+        agent_created_events = self.session.query(Event).filter(
             Event.org_id == self.organisation_id,
             Event.event_name == 'agent_created'
-        ).group_by(Event.agent_id).subquery()
-
-        other_tools = self.session.query(
-            Event.agent_id,
-            func.array_agg(distinct(Event.event_property['tool_name'].astext)).label('other_tools')
-        ).filter(
-            Event.org_id == self.organisation_id,
-            Event.event_name == 'tool_used',
-            Event.event_property['tool_name'].astext != formatted_tool_name
-        ).group_by(Event.agent_id).subquery()
-
-        result = self.session.query(
-            event_tool_used.c.agent_id,
-            event_tool_used.c.created_at,
-            event_tool_used.c.event_name,
-            event_run.c.tokens_consumed,
-            event_run.c.calls,
-            event_run_created.c.agent_execution_name,
-            event_agent_created.c.agent_name,
-            event_agent_created.c.model,
-            other_tools.c.other_tools
-        ).join(
-            event_run, event_tool_used.c.agent_id == event_run.c.agent_id
-        ).join(
-            event_run_created, event_tool_used.c.agent_id == event_run_created.c.agent_id
-        ).join(
-            event_agent_created, event_tool_used.c.agent_id == event_agent_created.c.agent_id
-        ).join(
-            other_tools, event_tool_used.c.agent_id == other_tools.c.agent_id, isouter=True
         ).all()
 
-        user_timezone = AgentConfiguration.get_agent_config_by_key_and_agent_id(session= self.session,key= 'user_timezone', agent_id= Event.agent_id)
-        if user_timezone.value is None:
-            user_timezone.value = 'GMT'
-
-        return [{
-            'agent_id': row.agent_id,
-            'created_at': row.created_at.astimezone(pytz.timezone(user_timezone.value)).strftime("%d %B %Y %H:%M"),
-            'event_name': row.event_name,
-            'tokens_consumed': row.tokens_consumed,
-            'calls': row.calls,
-            'agent_execution_name': row.agent_execution_name,
-            'agent_name': row.agent_name,
-            'model': row.model,
-            'other_tools': row.other_tools
-        } for row in result]
+        results = []
+
+        for tool_event in tool_events:
+            agent_execution_id = tool_event.event_property['agent_execution_id']
+
+            event_run = next((er for er in event_runs if er.agent_id == tool_event.agent_id and er.event_property['agent_execution_id'] == agent_execution_id), None)
+            agent_created_event = next((ace for ace in agent_created_events if ace.agent_id == tool_event.agent_id), None)
+            try:
+                user_timezone = AgentConfiguration.get_agent_config_by_key_and_agent_id(session=self.session, key='user_timezone', agent_id=tool_event.agent_id)
+                if user_timezone and user_timezone.value != 'None':
+                    tz = pytz.timezone(user_timezone.value)
+                else:
+                    tz = pytz.timezone('GMT')       
+            except AttributeError:
+                tz = pytz.timezone('GMT')
+
+            if event_run and agent_created_event:
+                actual_time = tool_event.created_at.astimezone(tz).strftime("%d %B %Y %H:%M")
+                other_tools_events = self.session.query(
+                    Event
+                ).filter(
+                    Event.org_id == self.organisation_id,
+                    Event.event_name == 'tool_used',
+                    Event.event_property['tool_name'].astext != formatted_tool_name,
+                    Event.agent_id == tool_event.agent_id, 
+                    Event.id.between(tool_event.id, event_run.id)
+                ).all()
+
+                other_tools = [ote.event_property['tool_name'] for ote in other_tools_events]
+
+                result_dict = {
+                    'created_at': actual_time,
+                    'agent_execution_id': agent_execution_id,
+                    'tokens_consumed': event_run.event_property['tokens_consumed'],
+                    'calls': event_run.event_property['calls'],
+                    'agent_execution_name': event_run.event_property['name'],
+                    'other_tools': other_tools,
+                    'agent_name': agent_created_event.event_property['agent_name'],
+                    'model': agent_created_event.event_property['model']
+                }
+
+                if agent_execution_id not in [i['agent_execution_id'] for i in results]:
+                    results.append(result_dict)
+
+        results = sorted(results, key=lambda x: datetime.strptime(x['created_at'], '%d %B %Y %H:%M'), reverse=True)
+
+        return results
+

superagi/controllers/agent.py

@@ -142,7 +142,8 @@ def create_agent_with_config(agent_with_config: AgentConfigInput,
     if agent_with_config.knowledge:
         knowledge_name = db.session.query(Knowledges.name).filter(Knowledges.id == agent_with_config.knowledge).first()[0]
         EventHandler(session=db.session).create_event('knowledge_picked', 
-                                                      {'knowledge_name': knowledge_name},
+                                                      {'knowledge_name': knowledge_name, 
+                                                        'agent_execution_id': execution.id},
                                                       db_agent.id, 
                                                       organisation.id if organisation else 0)
 

superagi/controllers/agent_execution.py

@@ -139,7 +139,8 @@ def create_agent_execution(agent_execution: AgentExecutionIn,
         knowledge_name = Knowledges.get_knowledge_from_id(db.session, int(agent_execution_knowledge.value)).name
         if knowledge_name is not None:
             EventHandler(session=db.session).create_event('knowledge_picked', 
-                                                        {'knowledge_name': knowledge_name},
+                                                        {'knowledge_name': knowledge_name, 
+                                                         'agent_execution_id': db_agent_execution.id},
                                                         agent_execution.agent_id, 
                                                         organisation.id if organisation else 0)
 
@@ -219,7 +220,8 @@ def create_agent_run(agent_execution: AgentRunIn, Authorize: AuthJWT = Depends(c
         knowledge_name = Knowledges.get_knowledge_from_id(db.session, int(agent_execution_knowledge.value)).name
         if knowledge_name is not None:
             EventHandler(session=db.session).create_event('knowledge_picked', 
-                                                        {'knowledge_name': knowledge_name},
+                                                        {'knowledge_name': knowledge_name, 
+                                                         'agent_execution_id': db_agent_execution.id},
                                                         agent_execution.agent_id, 
                                                         organisation.id if organisation else 0)
 

superagi/controllers/api/agent.py

@@ -157,7 +157,8 @@ def create_run(agent_id:int,agent_execution: AgentExecutionIn,api_key: str = Sec
         knowledge_name = Knowledges.get_knowledge_from_id(db.session, int(agent_execution_knowledge.value)).name
         if knowledge_name is not None:
             EventHandler(session=db.session).create_event('knowledge_picked', 
-                                                        {'knowledge_name': knowledge_name},
+                                                        {'knowledge_name': knowledge_name, 
+                                                         'agent_execution_id': db_agent_execution.id},
                                                         agent_id,
                                                         organisation.id if organisation else 0
                                                         )

superagi/jobs/scheduling_executor.py

@@ -67,7 +67,8 @@ def execute_scheduled_agent(self, agent_id: int, name: str):
             knowledge_name = Knowledges.get_knowledge_from_id(session, int(agent_execution_knowledge.value)).name
             if knowledge_name is not None:
                 EventHandler(session=session).create_event('knowledge_picked', 
-                                                        {'knowledge_name': knowledge_name},
+                                                        {'knowledge_name': knowledge_name, 
+                                                         'agent_execution_id': db_agent_execution.id},
                                                         agent_id, 
                                                         organisation.id if organisation else 0)
         session.commit()