Merge pull request #772 from TransformerOptimus/fixing_apm

Fixing apm queries

gui/pages/Content/APM/ApmDashboard.js

@@ -252,7 +252,7 @@ export default function ApmDashboard() {
                                                 <td className="table_data text_align_right" style={{width:'12%'}}>{run.runs_completed?(run.total_tokens/run.runs_completed).toFixed(1) : '-'}</td>
                                                 <td className="table_data text_align_right" style={{width:'20%'}}>
                                                     {run.tools_used && run.tools_used.slice(0, 3).map((tool,index) => (
-                                                        <div className="tools_used">{tool}</div>
+                                                        <div key={index} className="tools_used">{tool}</div>
                                                     ))}
                                                     {run.tools_used && run.tools_used.length > 3 &&
                                                         <div className="tools_used_tooltip"
@@ -277,7 +277,7 @@ export default function ApmDashboard() {
                                     <img src="/images/no_permissions.svg" width={190} height={74} alt="No Data"/>
                                     <span className="text_12 color_white mt_6">No active runs found</span>
                                 </div> : activeRuns.map((run,index) => (
-                                    <div className="active_runs">
+                                    <div key={index} className="active_runs">
                                         <span className="text_14">{run.name}</span>
                                         <div style={{display:'inline-flex',alignItems:'center'}}><span className="text_12 mt_6">{run.agent_name}  ·  <Image width={12} height={12} src="/images/schedule.svg" alt="schedule-icon" /> {formatTime(run.created_at)}</span></div>
                                     </div>

gui/pages/Content/APM/BarGraph.js

@@ -69,4 +69,6 @@ export const BarGraph = ({ data, type, color }) => {
             <div ref={chartRef} style={{ width: '100%', height: '100%' }}></div>
         </div>
     );
-}
+}
+
+export default BarGraph;

gui/pages/Content/Agents/ResourceList.js

@@ -51,6 +51,7 @@ export default function ResourceList({ files, channel, runs }) {
 
                         {selectedRun === filesRun.run && (
                             <div className={styles.resources} style={{padding: '2px 8px'}}>
+                              {/* eslint-disable-next-line react/jsx-key */}
                               {filesRun.files.map((file, index) => <File file={file} index={index} />)}
                             </div>
                         )}
@@ -61,6 +62,7 @@ export default function ResourceList({ files, channel, runs }) {
 
         {channel === 'input' &&
             <div className={styles.resources}>
+              {/* eslint-disable-next-line react/jsx-key */}
               {files.map((file, index) => <File file={file} index={index} />)}
             </div>}
       </div>

superagi/apm/analytics_helper.py

@@ -3,31 +3,32 @@
 from sqlalchemy.orm.query import Query
 from superagi.models.events import Event
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy import text, func, Integer
+from sqlalchemy import text, func, Integer, and_
 from collections import defaultdict
 import logging
 
 class AnalyticsHelper:
 
-    def __init__(self, session: Session):
+    def __init__(self, session: Session, organisation_id: int):
         self.session = session
+        self.organisation_id = organisation_id
 
     def calculate_run_completed_metrics(self) -> Dict[str, Dict[str, Union[int, List[Dict[str, int]]]]]:
 
         agent_model_query = self.session.query(
             Event.event_property['model'].label('model'),
             Event.agent_id
-        ).filter_by(event_name="agent_created").subquery()
+        ).filter_by(event_name="agent_created", org_id=self.organisation_id).subquery()
 
         agent_runs_query = self.session.query(
             agent_model_query.c.model,
             func.count(Event.id).label('runs')
-        ).join(Event, Event.agent_id == agent_model_query.c.agent_id).filter(Event.event_name.in_(['run_completed', 'run_iteration_limit_crossed'])).group_by(agent_model_query.c.model).subquery()
+        ).join(Event, and_(Event.agent_id == agent_model_query.c.agent_id, Event.org_id == self.organisation_id)).filter(Event.event_name.in_(['run_completed', 'run_iteration_limit_crossed'])).group_by(agent_model_query.c.model).subquery()
 
         agent_tokens_query = self.session.query(
             agent_model_query.c.model,
             func.sum(text("(event_property->>'tokens_consumed')::int")).label('tokens')
-        ).join(Event, Event.agent_id == agent_model_query.c.agent_id).filter(Event.event_name.in_(['run_completed', 'run_iteration_limit_crossed'])).group_by(agent_model_query.c.model).subquery()
+        ).join(Event, and_(Event.agent_id == agent_model_query.c.agent_id, Event.org_id == self.organisation_id)).filter(Event.event_name.in_(['run_completed', 'run_iteration_limit_crossed'])).group_by(agent_model_query.c.model).subquery()
 
         agent_count_query = self.session.query(
             agent_model_query.c.model,
@@ -60,31 +61,31 @@ def fetch_agent_data(self) -> Dict[str, List[Dict[str, Any]]]:
             Event.agent_id,
             Event.event_property['agent_name'].label('agent_name'),
             Event.event_property['model'].label('model')
-        ).filter_by(event_name="agent_created").subquery()
+        ).filter_by(event_name="agent_created", org_id=self.organisation_id).subquery()
 
         run_subquery = self.session.query(
             Event.agent_id,
             func.sum(text("(event_property->>'tokens_consumed')::int")).label('total_tokens'),
             func.sum(text("(event_property->>'calls')::int")).label('total_calls'),
             func.count(Event.id).label('runs_completed'),
-        ).filter(Event.event_name.in_(['run_completed', 'run_iteration_limit_crossed'])).group_by(Event.agent_id).subquery()
+        ).filter(and_(Event.event_name.in_(['run_completed', 'run_iteration_limit_crossed']), Event.org_id == self.organisation_id)).group_by(Event.agent_id).subquery()
 
         tool_subquery = self.session.query(
             Event.agent_id,
             func.array_agg(Event.event_property['tool_name'].distinct()).label('tools_used'),
-        ).filter_by(event_name="tool_used").group_by(Event.agent_id).subquery()
+        ).filter_by(event_name="tool_used", org_id=self.organisation_id).group_by(Event.agent_id).subquery()
 
         start_time_subquery = self.session.query(
             Event.agent_id,
-            (Event.event_property['agent_execution_id']).label('agent_execution_id'),
+            Event.event_property['agent_execution_id'].label('agent_execution_id'),
             func.min(func.extract('epoch', Event.created_at)).label('start_time')
-        ).filter_by(event_name="run_created").group_by(Event.agent_id, Event.event_property['agent_execution_id']).subquery()
+        ).filter_by(event_name="run_created", org_id=self.organisation_id).group_by(Event.agent_id, Event.event_property['agent_execution_id']).subquery()
 
         end_time_subquery = self.session.query(
             Event.agent_id,
-            (Event.event_property['agent_execution_id']).label('agent_execution_id'),
+            Event.event_property['agent_execution_id'].label('agent_execution_id'),
             func.max(func.extract('epoch', Event.created_at)).label('end_time')
-        ).filter(Event.event_name.in_(['run_completed', 'run_iteration_limit_crossed'])).group_by(Event.agent_id, Event.event_property['agent_execution_id']).subquery()
+        ).filter(and_(Event.event_name.in_(['run_completed', 'run_iteration_limit_crossed']), Event.org_id == self.organisation_id)).group_by(Event.agent_id, Event.event_property['agent_execution_id']).subquery()
 
         time_diff_subquery = self.session.query(
             start_time_subquery.c.agent_id,
@@ -128,13 +129,13 @@ def fetch_agent_runs(self, agent_id: int) -> List[Dict[str, int]]:
             Event.event_property['tokens_consumed'].label('tokens_consumed'),
             Event.event_property['calls'].label('calls'),
             Event.updated_at
-        ).filter(Event.event_name.in_(['run_completed','run_iteration_limit_crossed']), Event.agent_id==agent_id).subquery()
+        ).filter(Event.event_name.in_(['run_completed','run_iteration_limit_crossed']), Event.agent_id == agent_id, Event.org_id == self.organisation_id).subquery()
 
         created_subquery = self.session.query(
             Event.event_property['agent_execution_id'].label('created_agent_execution_id'),
             Event.event_property['agent_execution_name'].label('agent_execution_name'),
             Event.created_at
-        ).filter(Event.event_name=="run_created", Event.agent_id==agent_id).subquery()
+        ).filter(Event.event_name == "run_created", Event.agent_id == agent_id, Event.org_id == self.organisation_id).subquery()
 
         query = self.session.query(
             created_subquery.c.agent_execution_name,
@@ -163,20 +164,21 @@ def get_active_runs(self) -> List[Dict[str, str]]:
         end_event_subquery = self.session.query(
             Event.event_property['agent_execution_id'].label('agent_execution_id'),
         ).filter(
-            Event.event_name.in_(['run_completed', 'run_iteration_limit_crossed'])
+            Event.event_name.in_(['run_completed', 'run_iteration_limit_crossed']),
+            Event.org_id == self.organisation_id
         ).subquery()
 
         start_subquery = self.session.query(
             Event.event_property['agent_execution_id'].label('agent_execution_id'),
             Event.event_property['agent_execution_name'].label('agent_execution_name'),
             Event.created_at,
             Event.agent_id
-        ).filter_by(event_name="run_created").subquery()
+        ).filter_by(event_name="run_created", org_id = self.organisation_id).subquery()
 
         agent_created_subquery = self.session.query(
             Event.event_property['agent_name'].label('agent_name'),
             Event.agent_id
-        ).filter_by(event_name="agent_created").subquery()
+        ).filter_by(event_name="agent_created", org_id = self.organisation_id).subquery()
 
         query = self.session.query(
             start_subquery.c.agent_execution_name,

superagi/apm/tools_handler.py

@@ -5,15 +5,16 @@
 from superagi.models.events import Event
 class ToolsHandler:
 
-    def __init__(self, session: Session):
+    def __init__(self, session: Session, organisation_id: int):
         self.session = session
+        self.organisation_id = organisation_id
 
     def calculate_tool_usage(self) -> List[Dict[str, int]]:
         tool_usage = []
         tool_used_subquery = self.session.query(
             Event.event_property['tool_name'].label('tool_name'),
             Event.agent_id
-        ).filter_by(event_name="tool_used").subquery()
+        ).filter_by(event_name="tool_used", org_id=self.organisation_id).subquery()
 
         agent_count = self.session.query(
             tool_used_subquery.c.tool_name,

superagi/controllers/analytics.py

@@ -1,5 +1,5 @@
 from fastapi import APIRouter, Depends, HTTPException
-from superagi.helper.auth import check_auth
+from superagi.helper.auth import check_auth, get_user_organisation
 from superagi.apm.analytics_helper import AnalyticsHelper
 from superagi.apm.event_handler import EventHandler
 from superagi.apm.tools_handler import ToolsHandler
@@ -10,7 +10,7 @@
 router = APIRouter()
 
 @router.get("/metrics", status_code=200)
-def get_metrics(Authorize: AuthJWT = Depends(check_auth)):
+def get_metrics(organisation=Depends(get_user_organisation)):
     """
     Get the total tokens, total calls, and the number of run completed.
 
@@ -19,43 +19,43 @@ def get_metrics(Authorize: AuthJWT = Depends(check_auth)):
 
     """
     try:
-        return AnalyticsHelper(session=db.session).calculate_run_completed_metrics()
+        return AnalyticsHelper(session=db.session, organisation_id=organisation.id).calculate_run_completed_metrics()
     except Exception as e:
         logging.error(f"Error while calculating metrics: {str(e)}")
         raise HTTPException(status_code=500, detail="Internal Server Error")
 
 
 @router.get("/agents/all", status_code=200)
-def get_agents(Authorize: AuthJWT = Depends(check_auth)):
+def get_agents(organisation=Depends(get_user_organisation)):
     try:
-        return AnalyticsHelper(session=db.session).fetch_agent_data()
+        return AnalyticsHelper(session=db.session, organisation_id=organisation.id).fetch_agent_data()
     except Exception as e:
         logging.error(f"Error while fetching agent data: {str(e)}")
         raise HTTPException(status_code=500, detail="Internal Server Error")
 
 
 @router.get("/agents/{agent_id}", status_code=200)
-def get_agent_runs(agent_id: int, Authorize: AuthJWT = Depends(check_auth)):
+def get_agent_runs(agent_id: int, organisation=Depends(get_user_organisation)):
     try:
-        return AnalyticsHelper(session=db.session).fetch_agent_runs(agent_id)
+        return AnalyticsHelper(session=db.session, organisation_id=organisation.id).fetch_agent_runs(agent_id)
     except Exception as e:
         logging.error(f"Error while fetching agent runs: {str(e)}")
         raise HTTPException(status_code=500, detail="Internal Server Error")
 
 
 @router.get("/runs/active", status_code=200)
-def get_active_runs(Authorize: AuthJWT = Depends(check_auth)):
+def get_active_runs(organisation=Depends(get_user_organisation)):
     try:
-        return AnalyticsHelper(session=db.session).get_active_runs()
+        return AnalyticsHelper(session=db.session, organisation_id=organisation.id).get_active_runs()
     except Exception as e:
         logging.error(f"Error while getting active runs: {str(e)}")
         raise HTTPException(status_code=500, detail="Internal Server Error")
 
 
 @router.get("/tools/used", status_code=200)
-def get_tools_used(Authorize: AuthJWT = Depends(check_auth)):
+def get_tools_used(organisation=Depends(get_user_organisation)):
     try:
-        return ToolsHandler(session=db.session).calculate_tool_usage()
+        return ToolsHandler(session=db.session, organisation_id=organisation.id).calculate_tool_usage()
     except Exception as e:
         logging.error(f"Error while calculating tool usage: {str(e)}")
         raise HTTPException(status_code=500, detail="Internal Server Error")

tests/unit_tests/apm/test_analytics_helper.py

@@ -1,32 +1,41 @@
 import pytest
-from superagi.models.events import Event
-from superagi.apm.analytics_helper import AnalyticsHelper
 from unittest.mock import MagicMock
+from superagi.apm.analytics_helper import AnalyticsHelper
+from sqlalchemy.orm import Session
 
 @pytest.fixture
 def mock_session():
-    return MagicMock()
+    return MagicMock(spec=Session)
 
 @pytest.fixture
-def analytics_helper(mock_session):
-    return AnalyticsHelper(mock_session)
+def organisation_id():
+    return 1
+
+@pytest.fixture
+def analytics_helper(mock_session, organisation_id):
+    return AnalyticsHelper(mock_session, organisation_id)
 
 def test_calculate_run_completed_metrics(analytics_helper, mock_session):
-    mock_session.query().all.return_value = [MagicMock()]
+    analytics_helper.calculate_run_completed_metrics = MagicMock(return_value = {})
     result = analytics_helper.calculate_run_completed_metrics()
     assert isinstance(result, dict)
+    analytics_helper.calculate_run_completed_metrics.assert_called()
 
 def test_fetch_agent_data(analytics_helper, mock_session):
-    mock_session.query().all.return_value = [MagicMock()]
+    analytics_helper.fetch_agent_data = MagicMock(return_value = {})
     result = analytics_helper.fetch_agent_data()
     assert isinstance(result, dict)
+    analytics_helper.fetch_agent_data.assert_called()
 
 def test_fetch_agent_runs(analytics_helper, mock_session):
-    mock_session.query().all.return_value = [MagicMock()]
-    result = analytics_helper.fetch_agent_runs(1)
+    agent_id = 1
+    analytics_helper.fetch_agent_runs = MagicMock(return_value = [])
+    result = analytics_helper.fetch_agent_runs(agent_id)
     assert isinstance(result, list)
+    analytics_helper.fetch_agent_runs.assert_called_with(agent_id)
 
 def test_get_active_runs(analytics_helper, mock_session):
-    mock_session.query().all.return_value = [MagicMock()]
+    analytics_helper.get_active_runs = MagicMock(return_value = [])
     result = analytics_helper.get_active_runs()
-    assert isinstance(result, list)
+    assert isinstance(result, list)
+    analytics_helper.get_active_runs.assert_called()

tests/unit_tests/apm/test_tools_handler.py

@@ -1,17 +1,22 @@
 import pytest
 from unittest.mock import MagicMock
-
+from sqlalchemy.orm import Session
 from superagi.apm.tools_handler import ToolsHandler
 
 @pytest.fixture
 def mock_session():
-    return MagicMock()
+    return MagicMock(spec=Session)
+
+@pytest.fixture
+def organisation_id():
+    return 1
 
 @pytest.fixture
-def tools_handler(mock_session):
-    return ToolsHandler(mock_session)
+def tools_handler(mock_session, organisation_id):
+    return ToolsHandler(mock_session, organisation_id)
 
-def test_calculate_tool_usage(tools_handler, mock_session):
-    mock_session.query().all.return_value = [MagicMock()]
+def test_calculate_tool_usage(tools_handler):
+    tools_handler.calculate_tool_usage = MagicMock(return_value=[])
     result = tools_handler.calculate_tool_usage()
-    assert isinstance(result, list)
+    assert isinstance(result, list)
+    tools_handler.calculate_tool_usage.assert_called()