From fcc9431dee9cf8c6d4d30ebc71acdfee8621c521 Mon Sep 17 00:00:00 2001 From: sayan1101 Date: Sat, 26 Aug 2023 17:28:36 +0530 Subject: [PATCH 01/15] entity details --- superagi/apm/knowledge_handler.py | 120 +++++++++++++++++ superagi/apm/tools_handler.py | 121 +++++++++++++++++- superagi/controllers/agent.py | 28 ++-- superagi/controllers/agent_execution.py | 33 ++++- superagi/controllers/analytics.py | 42 ++++++ superagi/controllers/api/agent.py | 18 ++- superagi/jobs/scheduling_executor.py | 21 ++- .../unit_tests/apm/test_knowledge_handler.py | 84 ++++++++++++ tests/unit_tests/apm/test_tools_handler.py | 109 +++++++++++++++- 9 files changed, 543 insertions(+), 33 deletions(-) create mode 100644 superagi/apm/knowledge_handler.py create mode 100644 tests/unit_tests/apm/test_knowledge_handler.py diff --git a/superagi/apm/knowledge_handler.py b/superagi/apm/knowledge_handler.py new file mode 100644 index 000000000..74e3100bf --- /dev/null +++ b/superagi/apm/knowledge_handler.py @@ -0,0 +1,120 @@ +from sqlalchemy.orm import Session +from superagi.models.events import Event +from superagi.models.knowledges import Knowledges +from sqlalchemy import Integer, or_, label, case, and_ +from fastapi import HTTPException +from typing import List, Dict, Union, Any +from sqlalchemy.sql import func +from sqlalchemy.orm import aliased + + +class KnowledgeHandler: + def __init__(self, session: Session, organisation_id: int): + self.session = session + self.organisation_id = organisation_id + + + def get_knowledge_wise_usage(self) -> Dict[str, Dict[str, int]]: + + EventAlias = aliased(Event) + + query_data = self.session.query( + Event.event_property['knowledge_name'].label('knowledge_name'), + func.count(Event.agent_id.distinct()).label('knowledge_unique_agents') + ).filter( + Event.event_name == 'knowledge_picked', + Event.org_id == self.organisation_id + ).group_by( + Event.event_property['knowledge_name'] + ).all() + print(query_data) + + return { + record.knowledge_name: { + 'knowledge_unique_agents': record.knowledge_unique_agents, + 'knowledge_calls': self.session.query( + EventAlias + ).filter( + EventAlias.event_property['tool_name'].astext == 'knowledgesearch', + EventAlias.event_name == 'tool_used', + EventAlias.org_id == self.organisation_id, + EventAlias.agent_id.in_(self.session.query(Event.agent_id).filter( + Event.event_name == 'knowledge_picked', + Event.org_id == self.organisation_id, + Event.event_property['knowledge_name'].astext == record.knowledge_name + )) + ).count() + } for record in query_data + } + + + def get_knowledge_logs_by_name(self, knowledge_name: str) -> List[Dict[str, Union[str, int]]]: + + is_knowledge_valid = self.session.query(Knowledges).filter_by(name=knowledge_name).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( + 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() + + 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() + + event_run_created = self.session.query( + Event.agent_id, + label('agent_execution_name', func.max(Event.event_property['agent_execution_name'].astext)) + ).filter( + Event.org_id == self.organisation_id, + Event.event_name == 'run_created' + ).group_by(Event.agent_id).subquery() + + 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( + 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() + + return [{ + 'agent_id': row.agent_id, + 'created_at': row.created_at, + '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] \ No newline at end of file diff --git a/superagi/apm/tools_handler.py b/superagi/apm/tools_handler.py index 72048d529..13b42ebae 100644 --- a/superagi/apm/tools_handler.py +++ b/superagi/apm/tools_handler.py @@ -1,12 +1,13 @@ -from typing import List, Dict - -from sqlalchemy import func +from typing import List, Dict, Union +from sqlalchemy import func, distinct from sqlalchemy.orm import Session - +from sqlalchemy import Integer +from fastapi import HTTPException from superagi.models.events import Event from superagi.models.tool import Tool from superagi.models.toolkit import Toolkit - +from sqlalchemy import or_ +from sqlalchemy.sql import label class ToolsHandler: def __init__(self, session: Session, organisation_id: int): @@ -54,4 +55,112 @@ def calculate_tool_usage(self) -> List[Dict[str, int]]: 'toolkit': tool_and_toolkit.get(row.tool_name, None) } for row in result] - return tool_usage \ No newline at end of file + return tool_usage + + def get_tool_wise_usage(self) -> Dict[str, Dict[str, int]]: + + tool_used_events = self.session.query( + Event.event_property['tool_name'].label('tool_name'), + func.count(Event.id).label('tool_calls'), + func.count(distinct(Event.agent_id)).label('tool_unique_agents') + ).filter( + Event.event_name == 'tool_used', + Event.org_id == self.organisation_id + ).group_by( + Event.event_property['tool_name'] + ) + + tool_data = { + event.tool_name: { + 'tool_calls': event.tool_calls, + 'tool_unique_agents': event.tool_unique_agents + } for event in tool_used_events + } + + return tool_data + + def get_tool_logs_by_tool_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( + 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() + + 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() + + event_run_created = self.session.query( + Event.agent_id, + label('agent_execution_name', func.max(Event.event_property['agent_execution_name'].astext)) + ).filter( + Event.org_id == self.organisation_id, + Event.event_name == 'run_created' + ).group_by(Event.agent_id).subquery() + + 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( + 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(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() + + return [{ + 'agent_id': row.agent_id, + 'created_at': row.created_at, + '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] \ No newline at end of file diff --git a/superagi/controllers/agent.py b/superagi/controllers/agent.py index a912a722b..6a708cb6f 100644 --- a/superagi/controllers/agent.py +++ b/superagi/controllers/agent.py @@ -132,18 +132,26 @@ def create_agent_with_config(agent_with_config: AgentConfigInput, agent = db.session.query(Agent).filter(Agent.id == db_agent.id, ).first() organisation = agent.get_agent_organisation(db.session) - EventHandler(session=db.session).create_event('run_created', {'agent_execution_id': execution.id, - 'agent_execution_name': execution.name}, db_agent.id, - - organisation.id if organisation else 0), - EventHandler(session=db.session).create_event('agent_created', {'agent_name': agent_with_config.name, - - 'model': agent_with_config.model}, db_agent.id, - + + EventHandler(session=db.session).create_event('run_created', + {'agent_execution_id': execution.id, + 'agent_execution_name': execution.name}, + db_agent.id, + organisation.id if organisation else 0), + + 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}, + db_agent.id, + organisation.id if organisation else 0) + + EventHandler(session=db.session).create_event('agent_created', + {'agent_name': agent_with_config.name, + 'model': agent_with_config.model}, + db_agent.id, organisation.id if organisation else 0) - # execute_agent.delay(execution.id, datetime.now()) - db.session.commit() return { diff --git a/superagi/controllers/agent_execution.py b/superagi/controllers/agent_execution.py index 1a485562c..cfd73ad44 100644 --- a/superagi/controllers/agent_execution.py +++ b/superagi/controllers/agent_execution.py @@ -23,6 +23,7 @@ from superagi.apm.event_handler import EventHandler from superagi.controllers.tool import ToolOut from superagi.models.agent_config import AgentConfiguration +from superagi.models.knowledges import Knowledges router = APIRouter() @@ -121,9 +122,20 @@ def create_agent_execution(agent_execution: AgentExecutionIn, agent_execution_configs=agent_execution_configs) organisation = agent.get_agent_organisation(db.session) - EventHandler(session=db.session).create_event('run_created', {'agent_execution_id': db_agent_execution.id,'agent_execution_name':db_agent_execution.name}, - agent_execution.agent_id, organisation.id if organisation else 0) - + agent_execution_knowledge = db.session.query(AgentConfiguration).filter(AgentConfiguration.key == 'knowledge').filter(AgentConfiguration.agent_id == agent_execution.agent_id).first() + + EventHandler(session=db.session).create_event('run_created', + {'agent_execution_id': db_agent_execution.id, + 'agent_execution_name':db_agent_execution.name}, + agent_execution.agent_id, + organisation.id if organisation else 0) + if agent_execution_knowledge: + knowledge_name = db.session.query(Knowledges.name).filter(Knowledges.id == int(agent_execution_knowledge.value)).filter(Knowledges.organisation_id == organisation.id).first()[0] + EventHandler(session=db.session).create_event('knowledge_picked', + {'knowledge_name': knowledge_name}, + agent_execution.agent_id, + organisation.id if organisation else 0) + if db_agent_execution.status == "RUNNING": execute_agent.delay(db_agent_execution.id, datetime.now()) @@ -185,8 +197,19 @@ def create_agent_run(agent_execution: AgentRunIn, Authorize: AuthJWT = Depends(c agent_execution_configs = agent_execution_configs) organisation = agent.get_agent_organisation(db.session) - EventHandler(session=db.session).create_event('run_created', {'agent_execution_id': db_agent_execution.id,'agent_execution_name':db_agent_execution.name}, - agent_execution.agent_id, organisation.id if organisation else 0) + EventHandler(session=db.session).create_event('run_created', + {'agent_execution_id': db_agent_execution.id, + 'agent_execution_name':db_agent_execution.name}, + agent_execution.agent_id, + organisation.id if organisation else 0) + agent_execution_knowledge = db.session.query(AgentConfiguration).filter(AgentConfiguration.key == 'knowledge').filter(AgentConfiguration.agent_id == agent_execution.agent_id).first() + + if agent_execution_knowledge: + knowledge_name = db.session.query(Knowledges.name).filter(Knowledges.id == int(agent_execution_knowledge.value)).filter(Knowledges.organisation_id == organisation.id).first()[0] + EventHandler(session=db.session).create_event('knowledge_picked', + {'knowledge_name': knowledge_name}, + agent_execution.agent_id, + organisation.id if organisation else 0) if db_agent_execution.status == "RUNNING": execute_agent.delay(db_agent_execution.id, datetime.now()) diff --git a/superagi/controllers/analytics.py b/superagi/controllers/analytics.py index 5dbfec7d9..181158b4d 100644 --- a/superagi/controllers/analytics.py +++ b/superagi/controllers/analytics.py @@ -3,6 +3,7 @@ from superagi.apm.analytics_helper import AnalyticsHelper from superagi.apm.event_handler import EventHandler from superagi.apm.tools_handler import ToolsHandler +from superagi.apm.knowledge_handler import KnowledgeHandler from fastapi_jwt_auth import AuthJWT from fastapi_sqlalchemy import db import logging @@ -59,3 +60,44 @@ def get_tools_used(organisation=Depends(get_user_organisation)): except Exception as e: logging.error(f"Error while calculating tool usage: {str(e)}") raise HTTPException(status_code=500, detail="Internal Server Error") + + +@router.get("/tools/usage", status_code=200) +def get_tool_usage(organisation=Depends(get_user_organisation)): + try: + return ToolsHandler(session=db.session, organisation_id=organisation.id).get_tool_wise_usage() + except Exception as e: + logging.error(f"Error while getting the tools' usage details: {str(e)}") + raise HTTPException(status_code=500, detail="Internal Server Error") + + +@router.get("/knowledge/usage", status_code=200) +def get_knowledge_usage(organisation=Depends(get_user_organisation)): + try: + return KnowledgeHandler(session=db.session, organisation_id=organisation.id).get_knowledge_wise_usage() + except Exception as e: + logging.error(f"Error while getting the knowledge usage details: {str(e)}") + raise HTTPException(status_code=500, detail="Internal Server Error") + + +@router.get("/tools/{tool_name}/logs", status_code=200) +def get_tool_logs(tool_name: str, organisation=Depends(get_user_organisation)): + try: + return ToolsHandler(session=db.session, organisation_id=organisation.id).get_tool_logs_by_tool_name(tool_name) + except Exception as e: + logging.error(f"Error while getting tool log details: {str(e)}") + if e.status_code: + raise HTTPException(status_code=e.status_code, detail=e.detail) + else: + raise HTTPException(status_code=500, detail="Internal Server Error") + +@router.get("/knowledge/{knowledge_name}/logs", status_code=200) +def get_knowledge_logs(knowledge_name: str, organisation=Depends(get_user_organisation)): + try: + return KnowledgeHandler(session=db.session, organisation_id=organisation.id).get_knowledge_logs_by_name(knowledge_name) + except Exception as e: + logging.error(f"Error while getting tool log details: {str(e)}") + if e.status_code: + raise HTTPException(status_code=e.status_code, detail=e.detail) + else: + raise HTTPException(status_code=500, detail="Internal Server Error") \ No newline at end of file diff --git a/superagi/controllers/api/agent.py b/superagi/controllers/api/agent.py index 9fc397cca..bd9ad9deb 100644 --- a/superagi/controllers/api/agent.py +++ b/superagi/controllers/api/agent.py @@ -14,6 +14,7 @@ from superagi.models.workflows.agent_workflow import AgentWorkflow from superagi.models.agent_execution import AgentExecution from superagi.models.organisation import Organisation +from superagi.models.knowledges import Knowledges from superagi.models.resource import Resource from superagi.controllers.types.agent_with_config import AgentConfigExtInput,AgentConfigUpdateExtInput from superagi.models.workflows.iteration_workflow import IterationWorkflow @@ -144,8 +145,21 @@ def create_run(agent_id:int,agent_execution: AgentExecutionIn,api_key: str = Sec if agent_execution_configs != {}: AgentExecutionConfiguration.add_or_update_agent_execution_config(session=db.session, execution=db_agent_execution, agent_execution_configs=agent_execution_configs) - EventHandler(session=db.session).create_event('run_created', {'agent_execution_id': db_agent_execution.id,'agent_execution_name':db_agent_execution.name}, - agent_id, organisation.id if organisation else 0) + EventHandler(session=db.session).create_event('run_created', + {'agent_execution_id': db_agent_execution.id, + 'agent_execution_name':db_agent_execution.name + }, + agent_id, + organisation.id if organisation else 0) + + agent_execution_knowledge = db.session.query(AgentConfiguration).filter(AgentConfiguration.key == 'knowledge').filter(AgentConfiguration.agent_id == agent_id).first() + if agent_execution_knowledge: + knowledge_name = db.session.query(Knowledges.name).filter(Knowledges.id == int(agent_execution_knowledge.value)).filter(Knowledges.organisation_id == organisation.id).first()[0] + EventHandler(session=db.session).create_event('knowledge_picked', + {'knowledge_name': knowledge_name}, + agent_id, + organisation.id if organisation else 0 + ) if db_agent_execution.status == "RUNNING": execute_agent.delay(db_agent_execution.id, datetime.now()) diff --git a/superagi/jobs/scheduling_executor.py b/superagi/jobs/scheduling_executor.py index b7e62db4e..3ab395212 100644 --- a/superagi/jobs/scheduling_executor.py +++ b/superagi/jobs/scheduling_executor.py @@ -1,9 +1,7 @@ -import ast from datetime import datetime from fastapi import HTTPException from sqlalchemy.orm import sessionmaker -from superagi.models.tool import Tool from superagi.models.workflows.iteration_workflow import IterationWorkflow from superagi.worker import execute_agent @@ -13,7 +11,7 @@ from superagi.models.agent_execution import AgentExecution from superagi.models.agent_execution_config import AgentExecutionConfiguration from superagi.apm.event_handler import EventHandler - +from superagi.models.knowledges import Knowledges from superagi.models.db import connect_db @@ -56,12 +54,21 @@ def execute_scheduled_agent(self, agent_id: int, name: str): for agent_config in agent_configurations: agent_execution_config = AgentExecutionConfiguration(agent_execution_id=agent_execution_id, key=agent_config.key, value=agent_config.value) session.add(agent_execution_config) - - organisation = agent.get_agent_organisation(session) model = session.query(AgentConfiguration.value).filter(AgentConfiguration.agent_id == agent_id).filter(AgentConfiguration.key == 'model').first()[0] - EventHandler(session=session).create_event('run_created', {'agent_execution_id': db_agent_execution.id,'agent_execution_name':db_agent_execution.name}, agent_id, organisation.id if organisation else 0), - + # if knowledge_id: + EventHandler(session=session).create_event('run_created', + {'agent_execution_id': db_agent_execution.id, + 'agent_execution_name':db_agent_execution.name}, + agent_id, + organisation.id if organisation else 0) + agent_execution_knowledge = session.query(AgentConfiguration).filter(AgentConfiguration.key == 'knowledge').filter(AgentConfiguration.agent_id == agent_id).first() + if agent_execution_knowledge: + knowledge_name = session.query(Knowledges.name).filter(Knowledges.id == int(agent_execution_knowledge.value)).filter(Knowledges.organisation_id == organisation.id).first()[0] + EventHandler(session=session).create_event('knowledge_picked', + {'knowledge_name': knowledge_name}, + agent_id, + organisation.id if organisation else 0) session.commit() if db_agent_execution.status == "RUNNING": diff --git a/tests/unit_tests/apm/test_knowledge_handler.py b/tests/unit_tests/apm/test_knowledge_handler.py new file mode 100644 index 000000000..9f771f251 --- /dev/null +++ b/tests/unit_tests/apm/test_knowledge_handler.py @@ -0,0 +1,84 @@ +import pytest +from unittest.mock import MagicMock +from superagi.apm.knowledge_handler import KnowledgeHandler +from fastapi import HTTPException + +@pytest.fixture +def organisation_id(): + return 1 + +@pytest.fixture +def mock_session(): + return MagicMock() + +@pytest.fixture +def knowledge_handler(mock_session, organisation_id): + return KnowledgeHandler(mock_session, organisation_id) + +def test_get_knowledge_wise_usage(knowledge_handler, mock_session): + knowledge_handler.session = mock_session + mock_knowledge_event = MagicMock() + mock_knowledge_event.knowledge_name = 'Knowledge1' + mock_knowledge_event.knowledge_unique_agents = 5 + + mock_session.query.return_value.filter.return_value.group_by.return_value.all.return_value = [mock_knowledge_event] + mock_session.query.return_value.filter.return_value.count.return_value = 10 + + result = knowledge_handler.get_knowledge_wise_usage() + + assert isinstance(result, dict) + assert result == { + 'Knowledge1': { + 'knowledge_unique_agents': 5, + 'knowledge_calls': 10 + } + } + +def test_get_knowledge_logs_by_name(knowledge_handler, mock_session): + knowledge_name = 'knowledge1' + knowledge_handler.session = mock_session + + mock_knowledge = MagicMock() + mock_knowledge.name = 'knowledge1' + mock_session.query().filter_by().first.return_value = mock_knowledge + + result_obj = MagicMock() + result_obj.agent_id = 1 + result_obj.created_at = "2022-05-25" + result_obj.event_name = 'knowledge_picked' + result_obj.tokens_consumed = 10 + result_obj.calls = 5 + result_obj.agent_execution_name = 'Runner' + result_obj.agent_name = 'A1' + result_obj.model = 'M1' + + mock_subquery = MagicMock() + mock_subquery.c.agent_id = 1 + + mock_session.query().filter().group_by().subquery.return_value = mock_subquery + mock_session.query().join().join().join().all.return_value = [result_obj] + + result = knowledge_handler.get_knowledge_logs_by_name(knowledge_name) + + assert isinstance(result, list) + expected_result = [{ + 'agent_id': 1, + 'created_at': '2022-05-25', + 'event_name': 'knowledge_picked', + 'tokens_consumed': 10, + 'calls': 5, + 'agent_execution_name': 'Runner', + 'agent_name': 'A1', + 'model': 'M1' + }] + assert result == expected_result + +def test_get_knowledge_logs_by_name_knowledge_not_found(knowledge_handler, mock_session): + knowledge_name = "knowledge1" + + mock_session.query().filter_by().first.return_value = None + + with pytest.raises(HTTPException): + knowledge_handler.get_knowledge_logs_by_name(knowledge_name) + + assert mock_session.query().filter_by().first.called \ No newline at end of file diff --git a/tests/unit_tests/apm/test_tools_handler.py b/tests/unit_tests/apm/test_tools_handler.py index f58650cc7..64ba8d744 100644 --- a/tests/unit_tests/apm/test_tools_handler.py +++ b/tests/unit_tests/apm/test_tools_handler.py @@ -1,6 +1,6 @@ import pytest from unittest.mock import MagicMock - +from fastapi import HTTPException from superagi.apm.tools_handler import ToolsHandler from sqlalchemy.orm import Session @@ -17,6 +17,109 @@ 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()] + tool_used_subquery = MagicMock() + agent_count_subquery = MagicMock() + total_usage_subquery = MagicMock() + + tool_used_subquery.c.tool_name = 'Tool1' + tool_used_subquery.c.agent_id = 1 + + agent_count_subquery.c.tool_name = 'Tool1' + agent_count_subquery.c.unique_agents = 1 + + total_usage_subquery.c.tool_name = 'Tool1' + total_usage_subquery.c.total_usage = 5 + + tools_handler.get_tool_and_toolkit = MagicMock() + tools_handler.get_tool_and_toolkit.return_value = {'Tool1': 'Toolkit1'} + + mock_session.query().filter_by().subquery.return_value = tool_used_subquery + mock_session.query().group_by().subquery.return_value = agent_count_subquery + mock_session.query().group_by().subquery.return_value = total_usage_subquery + + result_obj = MagicMock() + result_obj.tool_name = 'Tool1' + result_obj.unique_agents = 1 + result_obj.total_usage = 5 + mock_session.query().join().all.return_value = [result_obj] + result = tools_handler.calculate_tool_usage() - assert isinstance(result, list) \ No newline at end of file + + assert isinstance(result, list) + + expected_output = [{'tool_name': 'Tool1', 'unique_agents': 1, 'total_usage': 5, 'toolkit': 'Toolkit1'}] + assert result == expected_output + +def test_get_tool_and_toolkit(tools_handler, mock_session): + result_obj = MagicMock() + result_obj.tool_name = 'tool 1' + result_obj.toolkit_name = 'toolkit 1' + + mock_session.query().join().all.return_value = [result_obj] + + output = tools_handler.get_tool_and_toolkit() + + assert isinstance(output, dict) + assert output == {'tool 1': 'toolkit 1'} + +def test_get_tool_wise_usage(tools_handler, mock_session): + tools_handler.session = mock_session + mock_tool_event = MagicMock() + mock_tool_event.tool_name = 'Tool1' + mock_tool_event.tool_calls = 10 + mock_tool_event.tool_unique_agents = 5 + + mock_session.query.return_value.filter.return_value.group_by.return_value = [mock_tool_event] + result = tools_handler.get_tool_wise_usage() + + assert isinstance(result, dict) + assert result == { + 'Tool1': { + 'tool_calls': 10, + 'tool_unique_agents': 5 + } + } + +def test_get_tool_logs_by_tool_name(tools_handler, mock_session): + tool_name = 'tool1' + + mock_tool = MagicMock() + mock_tool.name = "tool1" + mock_session.query().filter_by().first.return_value = mock_tool + + result_obj = MagicMock() + result_obj.agent_id = 1 + result_obj.created_at = "2022-05-25" + result_obj.event_name = 'tool_used' + result_obj.tokens_consumed = 10 + result_obj.calls = 5 + result_obj.agent_execution_name = 'Runner' + result_obj.agent_name = 'A1' + result_obj.model = 'M1' + result_obj.other_tools = ['tool2', 'tool3'] + + mock_session.query().join().join().join().join().all.return_value = [result_obj] + result = tools_handler.get_tool_logs_by_tool_name(tool_name) + + assert isinstance(result, list) + expected_result = [{ + 'agent_id': 1, + 'created_at': '2022-05-25', + 'event_name': 'tool_used', + 'tokens_consumed': 10, + 'calls': 5, + 'agent_execution_name': 'Runner', + 'agent_name': 'A1', + 'model': 'M1', + 'other_tools': ['tool2', 'tool3'] + }] + assert result == expected_result + +def test_get_tool_logs_by_tool_name_tool_not_found(tools_handler, mock_session): + tool_name = "tool1" + + mock_session.query().filter_by().first.return_value = None + with pytest.raises(HTTPException): + tools_handler.get_tool_logs_by_tool_name(tool_name) + + assert mock_session.query().filter_by().first.called \ No newline at end of file From f35d75fc29ca3378c5bd8ab0c038c061632adec4 Mon Sep 17 00:00:00 2001 From: sayan1101 Date: Sat, 26 Aug 2023 17:33:49 +0530 Subject: [PATCH 02/15] entity details --- superagi/apm/knowledge_handler.py | 2 +- superagi/apm/tools_handler.py | 2 +- superagi/controllers/analytics.py | 4 ++-- tests/unit_tests/apm/test_knowledge_handler.py | 8 ++++---- tests/unit_tests/apm/test_tools_handler.py | 8 ++++---- 5 files changed, 12 insertions(+), 12 deletions(-) diff --git a/superagi/apm/knowledge_handler.py b/superagi/apm/knowledge_handler.py index 74e3100bf..4436f7986 100644 --- a/superagi/apm/knowledge_handler.py +++ b/superagi/apm/knowledge_handler.py @@ -48,7 +48,7 @@ def get_knowledge_wise_usage(self) -> Dict[str, Dict[str, int]]: } - def get_knowledge_logs_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]]]: is_knowledge_valid = self.session.query(Knowledges).filter_by(name=knowledge_name).first() diff --git a/superagi/apm/tools_handler.py b/superagi/apm/tools_handler.py index 13b42ebae..77eb3cb79 100644 --- a/superagi/apm/tools_handler.py +++ b/superagi/apm/tools_handler.py @@ -79,7 +79,7 @@ def get_tool_wise_usage(self) -> Dict[str, Dict[str, int]]: return tool_data - def get_tool_logs_by_tool_name(self, tool_name: str) -> List[Dict[str, Union[str, int, List[str]]]]: + def get_tool_events_by_tool_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() diff --git a/superagi/controllers/analytics.py b/superagi/controllers/analytics.py index 181158b4d..427099b76 100644 --- a/superagi/controllers/analytics.py +++ b/superagi/controllers/analytics.py @@ -83,7 +83,7 @@ def get_knowledge_usage(organisation=Depends(get_user_organisation)): @router.get("/tools/{tool_name}/logs", status_code=200) def get_tool_logs(tool_name: str, organisation=Depends(get_user_organisation)): try: - return ToolsHandler(session=db.session, organisation_id=organisation.id).get_tool_logs_by_tool_name(tool_name) + return ToolsHandler(session=db.session, organisation_id=organisation.id).get_tool_events_by_tool_name(tool_name) except Exception as e: logging.error(f"Error while getting tool log details: {str(e)}") if e.status_code: @@ -94,7 +94,7 @@ def get_tool_logs(tool_name: str, organisation=Depends(get_user_organisation)): @router.get("/knowledge/{knowledge_name}/logs", status_code=200) def get_knowledge_logs(knowledge_name: str, organisation=Depends(get_user_organisation)): try: - return KnowledgeHandler(session=db.session, organisation_id=organisation.id).get_knowledge_logs_by_name(knowledge_name) + return KnowledgeHandler(session=db.session, organisation_id=organisation.id).get_knowledge_events_by_name(knowledge_name) except Exception as e: logging.error(f"Error while getting tool log details: {str(e)}") if e.status_code: diff --git a/tests/unit_tests/apm/test_knowledge_handler.py b/tests/unit_tests/apm/test_knowledge_handler.py index 9f771f251..0eb7acfec 100644 --- a/tests/unit_tests/apm/test_knowledge_handler.py +++ b/tests/unit_tests/apm/test_knowledge_handler.py @@ -34,7 +34,7 @@ def test_get_knowledge_wise_usage(knowledge_handler, mock_session): } } -def test_get_knowledge_logs_by_name(knowledge_handler, mock_session): +def test_get_knowledge_events_by_name(knowledge_handler, mock_session): knowledge_name = 'knowledge1' knowledge_handler.session = mock_session @@ -58,7 +58,7 @@ def test_get_knowledge_logs_by_name(knowledge_handler, mock_session): mock_session.query().filter().group_by().subquery.return_value = mock_subquery mock_session.query().join().join().join().all.return_value = [result_obj] - result = knowledge_handler.get_knowledge_logs_by_name(knowledge_name) + result = knowledge_handler.get_knowledge_events_by_name(knowledge_name) assert isinstance(result, list) expected_result = [{ @@ -73,12 +73,12 @@ def test_get_knowledge_logs_by_name(knowledge_handler, mock_session): }] assert result == expected_result -def test_get_knowledge_logs_by_name_knowledge_not_found(knowledge_handler, mock_session): +def test_get_knowledge_events_by_name_knowledge_not_found(knowledge_handler, mock_session): knowledge_name = "knowledge1" mock_session.query().filter_by().first.return_value = None with pytest.raises(HTTPException): - knowledge_handler.get_knowledge_logs_by_name(knowledge_name) + knowledge_handler.get_knowledge_events_by_name(knowledge_name) assert mock_session.query().filter_by().first.called \ No newline at end of file diff --git a/tests/unit_tests/apm/test_tools_handler.py b/tests/unit_tests/apm/test_tools_handler.py index 64ba8d744..e1b288394 100644 --- a/tests/unit_tests/apm/test_tools_handler.py +++ b/tests/unit_tests/apm/test_tools_handler.py @@ -80,7 +80,7 @@ def test_get_tool_wise_usage(tools_handler, mock_session): } } -def test_get_tool_logs_by_tool_name(tools_handler, mock_session): +def test_get_tool_events_by_tool_name(tools_handler, mock_session): tool_name = 'tool1' mock_tool = MagicMock() @@ -99,7 +99,7 @@ def test_get_tool_logs_by_tool_name(tools_handler, mock_session): result_obj.other_tools = ['tool2', 'tool3'] mock_session.query().join().join().join().join().all.return_value = [result_obj] - result = tools_handler.get_tool_logs_by_tool_name(tool_name) + result = tools_handler.get_tool_events_by_tool_name(tool_name) assert isinstance(result, list) expected_result = [{ @@ -115,11 +115,11 @@ def test_get_tool_logs_by_tool_name(tools_handler, mock_session): }] assert result == expected_result -def test_get_tool_logs_by_tool_name_tool_not_found(tools_handler, mock_session): +def test_get_tool_events_by_tool_name_tool_not_found(tools_handler, mock_session): tool_name = "tool1" mock_session.query().filter_by().first.return_value = None with pytest.raises(HTTPException): - tools_handler.get_tool_logs_by_tool_name(tool_name) + tools_handler.get_tool_events_by_tool_name(tool_name) assert mock_session.query().filter_by().first.called \ No newline at end of file From c7991e02b251ce7ff1fb52ec6b4ed82df24fd19e Mon Sep 17 00:00:00 2001 From: Sayan Samanta <139119661+sayan1101@users.noreply.github.com> Date: Sat, 26 Aug 2023 20:40:17 +0530 Subject: [PATCH 03/15] Update knowledge_handler.py --- superagi/apm/knowledge_handler.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/superagi/apm/knowledge_handler.py b/superagi/apm/knowledge_handler.py index 4436f7986..1f91861c7 100644 --- a/superagi/apm/knowledge_handler.py +++ b/superagi/apm/knowledge_handler.py @@ -27,7 +27,6 @@ def get_knowledge_wise_usage(self) -> Dict[str, Dict[str, int]]: ).group_by( Event.event_property['knowledge_name'] ).all() - print(query_data) return { record.knowledge_name: { @@ -117,4 +116,4 @@ def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Un 'agent_execution_name': row.agent_execution_name, 'agent_name': row.agent_name, 'model': row.model - } for row in result] \ No newline at end of file + } for row in result] From da200c93d0d9e9c260565c61097710b345e0d5e3 Mon Sep 17 00:00:00 2001 From: sayan1101 Date: Sat, 26 Aug 2023 23:11:35 +0530 Subject: [PATCH 04/15] removed print statement --- superagi/apm/knowledge_handler.py | 1 - 1 file changed, 1 deletion(-) diff --git a/superagi/apm/knowledge_handler.py b/superagi/apm/knowledge_handler.py index 4436f7986..8a1975b3e 100644 --- a/superagi/apm/knowledge_handler.py +++ b/superagi/apm/knowledge_handler.py @@ -27,7 +27,6 @@ def get_knowledge_wise_usage(self) -> Dict[str, Dict[str, int]]: ).group_by( Event.event_property['knowledge_name'] ).all() - print(query_data) return { record.knowledge_name: { From a2cf9319fdf276af4a591044d833cf3c608f29b4 Mon Sep 17 00:00:00 2001 From: Sayan Samanta <139119661+sayan1101@users.noreply.github.com> Date: Sun, 27 Aug 2023 10:55:17 +0530 Subject: [PATCH 05/15] Update knowledge_handler.py --- superagi/apm/knowledge_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superagi/apm/knowledge_handler.py b/superagi/apm/knowledge_handler.py index 1f91861c7..b7babd5c1 100644 --- a/superagi/apm/knowledge_handler.py +++ b/superagi/apm/knowledge_handler.py @@ -49,7 +49,7 @@ def get_knowledge_wise_usage(self) -> Dict[str, Dict[str, int]]: def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Union[str, int]]]: - is_knowledge_valid = self.session.query(Knowledges).filter_by(name=knowledge_name).first() + is_knowledge_valid = self.session.query(Knowledges.id).filter_by(name=knowledge_name).filter(Knowledges.organisation_id == self.organisation_id).first()[0] if not is_knowledge_valid: raise HTTPException(status_code=404, detail="Knowledge not found") From c98fbe62b8429e1110a7333de12902010095fdff Mon Sep 17 00:00:00 2001 From: Sayan Samanta <139119661+sayan1101@users.noreply.github.com> Date: Sun, 27 Aug 2023 11:08:34 +0530 Subject: [PATCH 06/15] Update analytics.py --- superagi/controllers/analytics.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/superagi/controllers/analytics.py b/superagi/controllers/analytics.py index 427099b76..adae683a2 100644 --- a/superagi/controllers/analytics.py +++ b/superagi/controllers/analytics.py @@ -86,7 +86,7 @@ def get_tool_logs(tool_name: str, organisation=Depends(get_user_organisation)): return ToolsHandler(session=db.session, organisation_id=organisation.id).get_tool_events_by_tool_name(tool_name) except Exception as e: logging.error(f"Error while getting tool log details: {str(e)}") - if e.status_code: + if hasattr(e, 'status_code'): raise HTTPException(status_code=e.status_code, detail=e.detail) else: raise HTTPException(status_code=500, detail="Internal Server Error") @@ -97,7 +97,7 @@ def get_knowledge_logs(knowledge_name: str, organisation=Depends(get_user_organi return KnowledgeHandler(session=db.session, organisation_id=organisation.id).get_knowledge_events_by_name(knowledge_name) except Exception as e: logging.error(f"Error while getting tool log details: {str(e)}") - if e.status_code: + if hasattr(e, 'status_code'): raise HTTPException(status_code=e.status_code, detail=e.detail) else: - raise HTTPException(status_code=500, detail="Internal Server Error") \ No newline at end of file + raise HTTPException(status_code=500, detail="Internal Server Error") From 6c82cc353a01adee31ed586d032e6aea8131aee4 Mon Sep 17 00:00:00 2001 From: Sayan Samanta <139119661+sayan1101@users.noreply.github.com> Date: Sun, 27 Aug 2023 11:14:43 +0530 Subject: [PATCH 07/15] Update knowledge_handler.py --- superagi/apm/knowledge_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superagi/apm/knowledge_handler.py b/superagi/apm/knowledge_handler.py index b7babd5c1..13997a637 100644 --- a/superagi/apm/knowledge_handler.py +++ b/superagi/apm/knowledge_handler.py @@ -49,7 +49,7 @@ def get_knowledge_wise_usage(self) -> Dict[str, Dict[str, int]]: def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Union[str, int]]]: - is_knowledge_valid = self.session.query(Knowledges.id).filter_by(name=knowledge_name).filter(Knowledges.organisation_id == self.organisation_id).first()[0] + 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") From 3ea87dca85514798ae3946154a6ec3757f111c95 Mon Sep 17 00:00:00 2001 From: Sayan Samanta <139119661+sayan1101@users.noreply.github.com> Date: Sun, 27 Aug 2023 11:20:11 +0530 Subject: [PATCH 08/15] Update test_knowledge_handler.py --- tests/unit_tests/apm/test_knowledge_handler.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/tests/unit_tests/apm/test_knowledge_handler.py b/tests/unit_tests/apm/test_knowledge_handler.py index 0eb7acfec..847ac5cd4 100644 --- a/tests/unit_tests/apm/test_knowledge_handler.py +++ b/tests/unit_tests/apm/test_knowledge_handler.py @@ -75,10 +75,14 @@ def test_get_knowledge_events_by_name(knowledge_handler, mock_session): def test_get_knowledge_events_by_name_knowledge_not_found(knowledge_handler, mock_session): knowledge_name = "knowledge1" - - mock_session.query().filter_by().first.return_value = None + not_found_message = 'Knowledge not found' + + mock_session.query().filter_by().filter().first.return_value = None - with pytest.raises(HTTPException): + try: knowledge_handler.get_knowledge_events_by_name(knowledge_name) - - assert mock_session.query().filter_by().first.called \ No newline at end of file + assert False, "Expected HTTPException has not been raised" + except HTTPException as e: + assert str(e.detail) == not_found_message, f"Expected {not_found_message}, got {e.detail}" + finally: + assert mock_session.query().filter_by().filter().first.called, "first() function not called" From 457d50daf1ee4b7654516e42de94b48ce39425dd Mon Sep 17 00:00:00 2001 From: Sayan Samanta <139119661+sayan1101@users.noreply.github.com> Date: Sun, 27 Aug 2023 16:34:14 +0530 Subject: [PATCH 09/15] Update analytics.py --- superagi/controllers/analytics.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/superagi/controllers/analytics.py b/superagi/controllers/analytics.py index adae683a2..e9fee1ad0 100644 --- a/superagi/controllers/analytics.py +++ b/superagi/controllers/analytics.py @@ -85,7 +85,7 @@ def get_tool_logs(tool_name: str, organisation=Depends(get_user_organisation)): try: return ToolsHandler(session=db.session, organisation_id=organisation.id).get_tool_events_by_tool_name(tool_name) except Exception as e: - logging.error(f"Error while getting tool log details: {str(e)}") + logging.error(f"Error while getting tool event details: {str(e)}") if hasattr(e, 'status_code'): raise HTTPException(status_code=e.status_code, detail=e.detail) else: @@ -96,7 +96,7 @@ def get_knowledge_logs(knowledge_name: str, organisation=Depends(get_user_organi try: return KnowledgeHandler(session=db.session, organisation_id=organisation.id).get_knowledge_events_by_name(knowledge_name) except Exception as e: - logging.error(f"Error while getting tool log details: {str(e)}") + logging.error(f"Error while getting knowledge event details: {str(e)}") if hasattr(e, 'status_code'): raise HTTPException(status_code=e.status_code, detail=e.detail) else: From 1e81c59291b609e17241d76bd718a7ef23f8e1d6 Mon Sep 17 00:00:00 2001 From: Sayan Samanta <139119661+sayan1101@users.noreply.github.com> Date: Sun, 27 Aug 2023 16:38:40 +0530 Subject: [PATCH 10/15] Update analytics.py From 7ace0d82b9a813985e97757a67a491e22752d933 Mon Sep 17 00:00:00 2001 From: sayan1101 Date: Sun, 27 Aug 2023 16:42:10 +0530 Subject: [PATCH 11/15] entity details fix --- superagi/apm/knowledge_handler.py | 3 +-- superagi/controllers/analytics.py | 8 ++++---- tests/unit_tests/apm/test_knowledge_handler.py | 14 +++++++++----- 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/superagi/apm/knowledge_handler.py b/superagi/apm/knowledge_handler.py index 8a1975b3e..12e6ca983 100644 --- a/superagi/apm/knowledge_handler.py +++ b/superagi/apm/knowledge_handler.py @@ -49,8 +49,7 @@ def get_knowledge_wise_usage(self) -> Dict[str, Dict[str, int]]: def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Union[str, int]]]: - is_knowledge_valid = self.session.query(Knowledges).filter_by(name=knowledge_name).first() - + 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") diff --git a/superagi/controllers/analytics.py b/superagi/controllers/analytics.py index 427099b76..3ef21868d 100644 --- a/superagi/controllers/analytics.py +++ b/superagi/controllers/analytics.py @@ -85,8 +85,8 @@ def get_tool_logs(tool_name: str, organisation=Depends(get_user_organisation)): try: return ToolsHandler(session=db.session, organisation_id=organisation.id).get_tool_events_by_tool_name(tool_name) except Exception as e: - logging.error(f"Error while getting tool log details: {str(e)}") - if e.status_code: + logging.error(f"Error while getting tool event details: {str(e)}") + if hasattr(e, 'status_code'): raise HTTPException(status_code=e.status_code, detail=e.detail) else: raise HTTPException(status_code=500, detail="Internal Server Error") @@ -96,8 +96,8 @@ def get_knowledge_logs(knowledge_name: str, organisation=Depends(get_user_organi try: return KnowledgeHandler(session=db.session, organisation_id=organisation.id).get_knowledge_events_by_name(knowledge_name) except Exception as e: - logging.error(f"Error while getting tool log details: {str(e)}") - if e.status_code: + logging.error(f"Error while getting knoelwdge event details: {str(e)}") + if hasattr(e, 'status_code'): raise HTTPException(status_code=e.status_code, detail=e.detail) else: raise HTTPException(status_code=500, detail="Internal Server Error") \ No newline at end of file diff --git a/tests/unit_tests/apm/test_knowledge_handler.py b/tests/unit_tests/apm/test_knowledge_handler.py index 0eb7acfec..d49ba24fb 100644 --- a/tests/unit_tests/apm/test_knowledge_handler.py +++ b/tests/unit_tests/apm/test_knowledge_handler.py @@ -75,10 +75,14 @@ def test_get_knowledge_events_by_name(knowledge_handler, mock_session): def test_get_knowledge_events_by_name_knowledge_not_found(knowledge_handler, mock_session): knowledge_name = "knowledge1" - - mock_session.query().filter_by().first.return_value = None + not_found_message = 'Knowledge not found' + + mock_session.query().filter_by().filter().first.return_value = None - with pytest.raises(HTTPException): + try: knowledge_handler.get_knowledge_events_by_name(knowledge_name) - - assert mock_session.query().filter_by().first.called \ No newline at end of file + assert False, "Expected HTTPException has not been raised" + except HTTPException as e: + assert str(e.detail) == not_found_message, f"Expected {not_found_message}, got {e.detail}" + finally: + assert mock_session.query().filter_by().filter().first.called, "first() function not called" \ No newline at end of file From e3b7da4b60fbc5850be7211e67efbd70743a6cc4 Mon Sep 17 00:00:00 2001 From: sayan1101 Date: Sun, 27 Aug 2023 22:16:17 +0530 Subject: [PATCH 12/15] entity details fix --- superagi/apm/knowledge_handler.py | 6 +++++- superagi/apm/tools_handler.py | 8 +++++--- superagi/controllers/analytics.py | 2 +- .../unit_tests/apm/test_knowledge_handler.py | 11 +++++++++-- tests/unit_tests/apm/test_tools_handler.py | 19 +++++++++++++------ 5 files changed, 33 insertions(+), 13 deletions(-) diff --git a/superagi/apm/knowledge_handler.py b/superagi/apm/knowledge_handler.py index 12e6ca983..9eba5e324 100644 --- a/superagi/apm/knowledge_handler.py +++ b/superagi/apm/knowledge_handler.py @@ -6,6 +6,8 @@ from typing import List, Dict, Union, Any from sqlalchemy.sql import func from sqlalchemy.orm import aliased +from superagi.models.agent_config import AgentConfiguration +import pytz class KnowledgeHandler: @@ -50,6 +52,7 @@ def get_knowledge_wise_usage(self) -> Dict[str, Dict[str, int]]: def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Union[str, int]]]: 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") @@ -106,9 +109,10 @@ def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Un event_agent_created, event_knowledge_picked.c.agent_id == event_agent_created.c.agent_id ).all() + user_timezone = self.session.query(AgentConfiguration).filter(AgentConfiguration.key == "user_timezone", AgentConfiguration.agent_id == Event.agent_id).first() return [{ 'agent_id': row.agent_id, - 'created_at': row.created_at, + '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, diff --git a/superagi/apm/tools_handler.py b/superagi/apm/tools_handler.py index 77eb3cb79..eb189bf97 100644 --- a/superagi/apm/tools_handler.py +++ b/superagi/apm/tools_handler.py @@ -8,6 +8,8 @@ from superagi.models.toolkit import Toolkit from sqlalchemy import or_ from sqlalchemy.sql import label +from superagi.models.agent_config import AgentConfiguration +import pytz class ToolsHandler: def __init__(self, session: Session, organisation_id: int): @@ -79,10 +81,9 @@ def get_tool_wise_usage(self) -> Dict[str, Dict[str, int]]: return tool_data - def get_tool_events_by_tool_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") @@ -153,9 +154,10 @@ def get_tool_events_by_tool_name(self, tool_name: str) -> List[Dict[str, Union[s other_tools, event_tool_used.c.agent_id == other_tools.c.agent_id, isouter=True ).all() + user_timezone = self.session.query(AgentConfiguration).filter(AgentConfiguration.key == "user_timezone", AgentConfiguration.agent_id == Event.agent_id).first() return [{ 'agent_id': row.agent_id, - 'created_at': row.created_at, + '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, diff --git a/superagi/controllers/analytics.py b/superagi/controllers/analytics.py index 3ef21868d..58709f7f5 100644 --- a/superagi/controllers/analytics.py +++ b/superagi/controllers/analytics.py @@ -83,7 +83,7 @@ def get_knowledge_usage(organisation=Depends(get_user_organisation)): @router.get("/tools/{tool_name}/logs", status_code=200) def get_tool_logs(tool_name: str, organisation=Depends(get_user_organisation)): try: - return ToolsHandler(session=db.session, organisation_id=organisation.id).get_tool_events_by_tool_name(tool_name) + return ToolsHandler(session=db.session, organisation_id=organisation.id).get_tool_events_by_name(tool_name) except Exception as e: logging.error(f"Error while getting tool event details: {str(e)}") if hasattr(e, 'status_code'): diff --git a/tests/unit_tests/apm/test_knowledge_handler.py b/tests/unit_tests/apm/test_knowledge_handler.py index d49ba24fb..46f5063af 100644 --- a/tests/unit_tests/apm/test_knowledge_handler.py +++ b/tests/unit_tests/apm/test_knowledge_handler.py @@ -2,6 +2,8 @@ from unittest.mock import MagicMock from superagi.apm.knowledge_handler import KnowledgeHandler from fastapi import HTTPException +from datetime import datetime +import pytz @pytest.fixture def organisation_id(): @@ -44,7 +46,7 @@ def test_get_knowledge_events_by_name(knowledge_handler, mock_session): result_obj = MagicMock() result_obj.agent_id = 1 - result_obj.created_at = "2022-05-25" + result_obj.created_at = datetime.now() result_obj.event_name = 'knowledge_picked' result_obj.tokens_consumed = 10 result_obj.calls = 5 @@ -58,12 +60,17 @@ def test_get_knowledge_events_by_name(knowledge_handler, mock_session): mock_session.query().filter().group_by().subquery.return_value = mock_subquery mock_session.query().join().join().join().all.return_value = [result_obj] + user_timezone = MagicMock() + mock_session.query().filter().first.return_value = user_timezone + user_timezone.value = 'America/New_York' + result = knowledge_handler.get_knowledge_events_by_name(knowledge_name) assert isinstance(result, list) + expected_result = [{ 'agent_id': 1, - 'created_at': '2022-05-25', + 'created_at': result_obj.created_at.astimezone(pytz.timezone(user_timezone.value)).strftime("%d %B %Y %H:%M"), 'event_name': 'knowledge_picked', 'tokens_consumed': 10, 'calls': 5, diff --git a/tests/unit_tests/apm/test_tools_handler.py b/tests/unit_tests/apm/test_tools_handler.py index e1b288394..4df85a09e 100644 --- a/tests/unit_tests/apm/test_tools_handler.py +++ b/tests/unit_tests/apm/test_tools_handler.py @@ -3,6 +3,8 @@ from fastapi import HTTPException from superagi.apm.tools_handler import ToolsHandler from sqlalchemy.orm import Session +from datetime import datetime +import pytz @pytest.fixture def organisation_id(): @@ -80,7 +82,7 @@ def test_get_tool_wise_usage(tools_handler, mock_session): } } -def test_get_tool_events_by_tool_name(tools_handler, mock_session): +def test_get_tool_events_by_name(tools_handler, mock_session): tool_name = 'tool1' mock_tool = MagicMock() @@ -89,7 +91,7 @@ def test_get_tool_events_by_tool_name(tools_handler, mock_session): result_obj = MagicMock() result_obj.agent_id = 1 - result_obj.created_at = "2022-05-25" + result_obj.created_at = datetime.now() # Set created_at to datetime not string result_obj.event_name = 'tool_used' result_obj.tokens_consumed = 10 result_obj.calls = 5 @@ -98,13 +100,18 @@ def test_get_tool_events_by_tool_name(tools_handler, mock_session): result_obj.model = 'M1' result_obj.other_tools = ['tool2', 'tool3'] + user_timezone = MagicMock() + user_timezone.value = 'UTC' + mock_session.query().filter().first.return_value = user_timezone + mock_session.query().join().join().join().join().all.return_value = [result_obj] - result = tools_handler.get_tool_events_by_tool_name(tool_name) + result = tools_handler.get_tool_events_by_name(tool_name) assert isinstance(result, list) + expected_result = [{ 'agent_id': 1, - 'created_at': '2022-05-25', + 'created_at': result_obj.created_at.astimezone(pytz.timezone(user_timezone.value)).strftime("%d %B %Y %H:%M"), 'event_name': 'tool_used', 'tokens_consumed': 10, 'calls': 5, @@ -115,11 +122,11 @@ def test_get_tool_events_by_tool_name(tools_handler, mock_session): }] assert result == expected_result -def test_get_tool_events_by_tool_name_tool_not_found(tools_handler, mock_session): +def test_get_tool_events_by_name_tool_not_found(tools_handler, mock_session): tool_name = "tool1" mock_session.query().filter_by().first.return_value = None with pytest.raises(HTTPException): - tools_handler.get_tool_events_by_tool_name(tool_name) + tools_handler.get_tool_events_by_name(tool_name) assert mock_session.query().filter_by().first.called \ No newline at end of file From 6ec17fe9870e7f6f8f84a49656b8ea1fd27d6d57 Mon Sep 17 00:00:00 2001 From: sayan1101 Date: Mon, 28 Aug 2023 12:29:13 +0530 Subject: [PATCH 13/15] entity fix --- superagi/apm/knowledge_handler.py | 9 +++++++-- superagi/apm/tools_handler.py | 6 +++++- superagi/controllers/agent_execution.py | 27 +++++++++++++------------ superagi/controllers/api/agent.py | 15 +++++++------- superagi/jobs/scheduling_executor.py | 13 ++++++------ superagi/models/agent_config.py | 9 +++++++++ 6 files changed, 50 insertions(+), 29 deletions(-) diff --git a/superagi/apm/knowledge_handler.py b/superagi/apm/knowledge_handler.py index 721af2f9b..3195af3bb 100644 --- a/superagi/apm/knowledge_handler.py +++ b/superagi/apm/knowledge_handler.py @@ -29,6 +29,7 @@ def get_knowledge_wise_usage(self) -> Dict[str, Dict[str, int]]: ).group_by( Event.event_property['knowledge_name'] ).all() + print(query_data) return { record.knowledge_name: { @@ -92,6 +93,7 @@ def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Un 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, @@ -109,7 +111,10 @@ def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Un event_agent_created, event_knowledge_picked.c.agent_id == event_agent_created.c.agent_id ).all() - user_timezone = self.session.query(AgentConfiguration).filter(AgentConfiguration.key == "user_timezone", AgentConfiguration.agent_id == Event.agent_id).first() + 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"), @@ -119,4 +124,4 @@ def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Un 'agent_execution_name': row.agent_execution_name, 'agent_name': row.agent_name, 'model': row.model - } for row in result] + } for row in result] \ No newline at end of file diff --git a/superagi/apm/tools_handler.py b/superagi/apm/tools_handler.py index eb189bf97..f7f852e89 100644 --- a/superagi/apm/tools_handler.py +++ b/superagi/apm/tools_handler.py @@ -84,6 +84,7 @@ def get_tool_wise_usage(self) -> Dict[str, Dict[str, int]]: 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") @@ -154,7 +155,10 @@ def get_tool_events_by_name(self, tool_name: str) -> List[Dict[str, Union[str, i other_tools, event_tool_used.c.agent_id == other_tools.c.agent_id, isouter=True ).all() - user_timezone = self.session.query(AgentConfiguration).filter(AgentConfiguration.key == "user_timezone", AgentConfiguration.agent_id == Event.agent_id).first() + 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"), diff --git a/superagi/controllers/agent_execution.py b/superagi/controllers/agent_execution.py index cfd73ad44..f8872f192 100644 --- a/superagi/controllers/agent_execution.py +++ b/superagi/controllers/agent_execution.py @@ -122,7 +122,7 @@ def create_agent_execution(agent_execution: AgentExecutionIn, agent_execution_configs=agent_execution_configs) organisation = agent.get_agent_organisation(db.session) - agent_execution_knowledge = db.session.query(AgentConfiguration).filter(AgentConfiguration.key == 'knowledge').filter(AgentConfiguration.agent_id == agent_execution.agent_id).first() + agent_execution_knowledge = AgentConfiguration.get_agent_config_by_key_and_agent_id(db.session, 'knowledge', agent_execution.agent_id) EventHandler(session=db.session).create_event('run_created', {'agent_execution_id': db_agent_execution.id, @@ -130,11 +130,12 @@ def create_agent_execution(agent_execution: AgentExecutionIn, agent_execution.agent_id, organisation.id if organisation else 0) if agent_execution_knowledge: - knowledge_name = db.session.query(Knowledges.name).filter(Knowledges.id == int(agent_execution_knowledge.value)).filter(Knowledges.organisation_id == organisation.id).first()[0] - EventHandler(session=db.session).create_event('knowledge_picked', - {'knowledge_name': knowledge_name}, - agent_execution.agent_id, - organisation.id if organisation else 0) + 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}, + agent_execution.agent_id, + organisation.id if organisation else 0) if db_agent_execution.status == "RUNNING": execute_agent.delay(db_agent_execution.id, datetime.now()) @@ -202,14 +203,14 @@ def create_agent_run(agent_execution: AgentRunIn, Authorize: AuthJWT = Depends(c 'agent_execution_name':db_agent_execution.name}, agent_execution.agent_id, organisation.id if organisation else 0) - agent_execution_knowledge = db.session.query(AgentConfiguration).filter(AgentConfiguration.key == 'knowledge').filter(AgentConfiguration.agent_id == agent_execution.agent_id).first() - + agent_execution_knowledge = AgentConfiguration.get_agent_config_by_key_and_agent_id(db.session, 'knowledge', agent_execution.agent_id) if agent_execution_knowledge: - knowledge_name = db.session.query(Knowledges.name).filter(Knowledges.id == int(agent_execution_knowledge.value)).filter(Knowledges.organisation_id == organisation.id).first()[0] - EventHandler(session=db.session).create_event('knowledge_picked', - {'knowledge_name': knowledge_name}, - agent_execution.agent_id, - organisation.id if organisation else 0) + 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}, + agent_execution.agent_id, + organisation.id if organisation else 0) if db_agent_execution.status == "RUNNING": execute_agent.delay(db_agent_execution.id, datetime.now()) diff --git a/superagi/controllers/api/agent.py b/superagi/controllers/api/agent.py index bd9ad9deb..a0e9a5e13 100644 --- a/superagi/controllers/api/agent.py +++ b/superagi/controllers/api/agent.py @@ -152,14 +152,15 @@ def create_run(agent_id:int,agent_execution: AgentExecutionIn,api_key: str = Sec agent_id, organisation.id if organisation else 0) - agent_execution_knowledge = db.session.query(AgentConfiguration).filter(AgentConfiguration.key == 'knowledge').filter(AgentConfiguration.agent_id == agent_id).first() + agent_execution_knowledge = AgentConfiguration.get_agent_config_by_key_and_agent_id(db.session, 'knowledge', agent_id) if agent_execution_knowledge: - knowledge_name = db.session.query(Knowledges.name).filter(Knowledges.id == int(agent_execution_knowledge.value)).filter(Knowledges.organisation_id == organisation.id).first()[0] - EventHandler(session=db.session).create_event('knowledge_picked', - {'knowledge_name': knowledge_name}, - agent_id, - organisation.id if organisation else 0 - ) + 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}, + agent_id, + organisation.id if organisation else 0 + ) if db_agent_execution.status == "RUNNING": execute_agent.delay(db_agent_execution.id, datetime.now()) diff --git a/superagi/jobs/scheduling_executor.py b/superagi/jobs/scheduling_executor.py index 3ab395212..f3051b0ec 100644 --- a/superagi/jobs/scheduling_executor.py +++ b/superagi/jobs/scheduling_executor.py @@ -62,13 +62,14 @@ def execute_scheduled_agent(self, agent_id: int, name: str): 'agent_execution_name':db_agent_execution.name}, agent_id, organisation.id if organisation else 0) - agent_execution_knowledge = session.query(AgentConfiguration).filter(AgentConfiguration.key == 'knowledge').filter(AgentConfiguration.agent_id == agent_id).first() + agent_execution_knowledge = AgentConfiguration.get_agent_config_by_key_and_agent_id(session, 'knowledge', agent_id) if agent_execution_knowledge: - knowledge_name = session.query(Knowledges.name).filter(Knowledges.id == int(agent_execution_knowledge.value)).filter(Knowledges.organisation_id == organisation.id).first()[0] - EventHandler(session=session).create_event('knowledge_picked', - {'knowledge_name': knowledge_name}, - agent_id, - organisation.id if organisation else 0) + 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}, + agent_id, + organisation.id if organisation else 0) session.commit() if db_agent_execution.status == "RUNNING": diff --git a/superagi/models/agent_config.py b/superagi/models/agent_config.py index 67b377da3..50d8d195b 100644 --- a/superagi/models/agent_config.py +++ b/superagi/models/agent_config.py @@ -115,3 +115,12 @@ def get_model_api_key(cls, session, agent_id: int, model: str): # if selected_model_source == ModelSourceType.Replicate: # return get_config("REPLICATE_API_TOKEN") # return get_config("OPENAI_API_KEY") +# + @classmethod + def get_agent_config_by_key_and_agent_id(cls, session, key: str, agent_id: int): + knowledge_config = session.query(AgentConfiguration).filter( + AgentConfiguration.agent_id == agent_id, + AgentConfiguration.key == key + ).first() + + return knowledge_config \ No newline at end of file From 25676abdcf4bf9979bcb8bf342b4aeec5c540c9a Mon Sep 17 00:00:00 2001 From: sayan1101 Date: Mon, 28 Aug 2023 23:02:22 +0530 Subject: [PATCH 14/15] entity fix --- superagi/apm/knowledge_handler.py | 28 +++++++++++------- superagi/apm/tools_handler.py | 25 ++++++++++------ superagi/controllers/agent_execution.py | 4 +-- superagi/controllers/analytics.py | 28 ++++++++++-------- superagi/controllers/api/agent.py | 2 +- superagi/jobs/scheduling_executor.py | 2 +- superagi/models/agent_config.py | 6 ++-- .../unit_tests/apm/test_knowledge_handler.py | 23 ++++++++++----- tests/unit_tests/apm/test_tools_handler.py | 29 +++++++++++++------ 9 files changed, 91 insertions(+), 56 deletions(-) diff --git a/superagi/apm/knowledge_handler.py b/superagi/apm/knowledge_handler.py index 3195af3bb..54abcd286 100644 --- a/superagi/apm/knowledge_handler.py +++ b/superagi/apm/knowledge_handler.py @@ -16,24 +16,29 @@ def __init__(self, session: Session, organisation_id: int): self.organisation_id = organisation_id - def get_knowledge_wise_usage(self) -> Dict[str, Dict[str, int]]: + def get_knowledge_usage_by_name(self, knowledge_name: str) -> Dict[str, Dict[str, int]]: + 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") EventAlias = aliased(Event) - query_data = self.session.query( + knowledge_used_event = self.session.query( Event.event_property['knowledge_name'].label('knowledge_name'), func.count(Event.agent_id.distinct()).label('knowledge_unique_agents') ).filter( Event.event_name == 'knowledge_picked', - Event.org_id == self.organisation_id + Event.org_id == self.organisation_id, + Event.event_property['knowledge_name'].astext == knowledge_name ).group_by( Event.event_property['knowledge_name'] - ).all() - print(query_data) + ).first() - return { - record.knowledge_name: { - 'knowledge_unique_agents': record.knowledge_unique_agents, + if knowledge_used_event is None: + return {} + + knowledge_data = { + 'knowledge_unique_agents': knowledge_used_event.knowledge_unique_agents, 'knowledge_calls': self.session.query( EventAlias ).filter( @@ -43,11 +48,12 @@ def get_knowledge_wise_usage(self) -> Dict[str, Dict[str, int]]: EventAlias.agent_id.in_(self.session.query(Event.agent_id).filter( Event.event_name == 'knowledge_picked', Event.org_id == self.organisation_id, - Event.event_property['knowledge_name'].astext == record.knowledge_name + Event.event_property['knowledge_name'].astext == knowledge_name )) ).count() - } for record in query_data - } + } + + return knowledge_data def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Union[str, int]]]: diff --git a/superagi/apm/tools_handler.py b/superagi/apm/tools_handler.py index f7f852e89..11e21c882 100644 --- a/superagi/apm/tools_handler.py +++ b/superagi/apm/tools_handler.py @@ -59,25 +59,32 @@ def calculate_tool_usage(self) -> List[Dict[str, int]]: return tool_usage - def get_tool_wise_usage(self) -> Dict[str, Dict[str, int]]: + def get_tool_usage_by_name(self, tool_name: str) -> Dict[str, Dict[str, int]]: + 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(" ", "") - tool_used_events = self.session.query( + tool_used_event = self.session.query( Event.event_property['tool_name'].label('tool_name'), func.count(Event.id).label('tool_calls'), func.count(distinct(Event.agent_id)).label('tool_unique_agents') ).filter( Event.event_name == 'tool_used', - Event.org_id == self.organisation_id + Event.org_id == self.organisation_id, + Event.event_property['tool_name'].astext == formatted_tool_name ).group_by( Event.event_property['tool_name'] - ) + ).first() + + if tool_used_event is None: + return {} tool_data = { - event.tool_name: { - 'tool_calls': event.tool_calls, - 'tool_unique_agents': event.tool_unique_agents - } for event in tool_used_events - } + 'tool_calls': tool_used_event.tool_calls, + 'tool_unique_agents': tool_used_event.tool_unique_agents + } return tool_data diff --git a/superagi/controllers/agent_execution.py b/superagi/controllers/agent_execution.py index f8872f192..8da797dc8 100644 --- a/superagi/controllers/agent_execution.py +++ b/superagi/controllers/agent_execution.py @@ -122,7 +122,7 @@ def create_agent_execution(agent_execution: AgentExecutionIn, agent_execution_configs=agent_execution_configs) organisation = agent.get_agent_organisation(db.session) - agent_execution_knowledge = AgentConfiguration.get_agent_config_by_key_and_agent_id(db.session, 'knowledge', agent_execution.agent_id) + agent_execution_knowledge = AgentConfiguration.get_agent_config_by_key_and_agent_id(session= db.session, key= 'knowledge', agent_id= agent_execution.agent_id) EventHandler(session=db.session).create_event('run_created', {'agent_execution_id': db_agent_execution.id, @@ -203,7 +203,7 @@ def create_agent_run(agent_execution: AgentRunIn, Authorize: AuthJWT = Depends(c 'agent_execution_name':db_agent_execution.name}, agent_execution.agent_id, organisation.id if organisation else 0) - agent_execution_knowledge = AgentConfiguration.get_agent_config_by_key_and_agent_id(db.session, 'knowledge', agent_execution.agent_id) + agent_execution_knowledge = AgentConfiguration.get_agent_config_by_key_and_agent_id(session= db.session, key= 'knowledge', agent_id= agent_execution.agent_id) if agent_execution_knowledge: knowledge_name = Knowledges.get_knowledge_from_id(db.session, int(agent_execution_knowledge.value)).name if knowledge_name is not None: diff --git a/superagi/controllers/analytics.py b/superagi/controllers/analytics.py index 6d8f76281..c2c78f764 100644 --- a/superagi/controllers/analytics.py +++ b/superagi/controllers/analytics.py @@ -62,22 +62,26 @@ def get_tools_used(organisation=Depends(get_user_organisation)): raise HTTPException(status_code=500, detail="Internal Server Error") -@router.get("/tools/usage", status_code=200) -def get_tool_usage(organisation=Depends(get_user_organisation)): +@router.get("/tools/{tool_name}/usage", status_code=200) +def get_tool_usage(tool_name: str, organisation=Depends(get_user_organisation)): try: - return ToolsHandler(session=db.session, organisation_id=organisation.id).get_tool_wise_usage() + return ToolsHandler(session=db.session, organisation_id=organisation.id).get_tool_usage_by_name(tool_name) except Exception as e: - logging.error(f"Error while getting the tools' usage details: {str(e)}") - raise HTTPException(status_code=500, detail="Internal Server Error") + if hasattr(e, 'status_code'): + raise HTTPException(status_code=e.status_code, detail=e.detail) + else: + raise HTTPException(status_code=500, detail="Internal Server Error") -@router.get("/knowledge/usage", status_code=200) -def get_knowledge_usage(organisation=Depends(get_user_organisation)): +@router.get("/knowledge/{knowledge_name}/usage", status_code=200) +def get_knowledge_usage(knowledge_name:str, organisation=Depends(get_user_organisation)): try: - return KnowledgeHandler(session=db.session, organisation_id=organisation.id).get_knowledge_wise_usage() + return KnowledgeHandler(session=db.session, organisation_id=organisation.id).get_knowledge_usage_by_name(knowledge_name) except Exception as e: - logging.error(f"Error while getting the knowledge usage details: {str(e)}") - raise HTTPException(status_code=500, detail="Internal Server Error") + if hasattr(e, 'status_code'): + raise HTTPException(status_code=e.status_code, detail=e.detail) + else: + raise HTTPException(status_code=500, detail="Internal Server Error") @router.get("/tools/{tool_name}/logs", status_code=200) @@ -89,7 +93,7 @@ def get_tool_logs(tool_name: str, organisation=Depends(get_user_organisation)): if hasattr(e, 'status_code'): raise HTTPException(status_code=e.status_code, detail=e.detail) else: - raise HTTPException(status_code=500, detail="Internal Server Error") + raise HTTPException(status_code=500, detail="Internal Server Error") @router.get("/knowledge/{knowledge_name}/logs", status_code=200) def get_knowledge_logs(knowledge_name: str, organisation=Depends(get_user_organisation)): @@ -100,4 +104,4 @@ def get_knowledge_logs(knowledge_name: str, organisation=Depends(get_user_organi if hasattr(e, 'status_code'): raise HTTPException(status_code=e.status_code, detail=e.detail) else: - raise HTTPException(status_code=500, detail="Internal Server Error") + raise HTTPException(status_code=500, detail="Internal Server Error") \ No newline at end of file diff --git a/superagi/controllers/api/agent.py b/superagi/controllers/api/agent.py index 074bd2faf..b1ef95140 100644 --- a/superagi/controllers/api/agent.py +++ b/superagi/controllers/api/agent.py @@ -152,7 +152,7 @@ def create_run(agent_id:int,agent_execution: AgentExecutionIn,api_key: str = Sec agent_id, organisation.id if organisation else 0) - agent_execution_knowledge = AgentConfiguration.get_agent_config_by_key_and_agent_id(db.session, 'knowledge', agent_id) + agent_execution_knowledge = AgentConfiguration.get_agent_config_by_key_and_agent_id(session= db.session, key= 'knowledge', agent_id= agent_id) if agent_execution_knowledge: knowledge_name = Knowledges.get_knowledge_from_id(db.session, int(agent_execution_knowledge.value)).name if knowledge_name is not None: diff --git a/superagi/jobs/scheduling_executor.py b/superagi/jobs/scheduling_executor.py index f3051b0ec..31039839e 100644 --- a/superagi/jobs/scheduling_executor.py +++ b/superagi/jobs/scheduling_executor.py @@ -62,7 +62,7 @@ def execute_scheduled_agent(self, agent_id: int, name: str): 'agent_execution_name':db_agent_execution.name}, agent_id, organisation.id if organisation else 0) - agent_execution_knowledge = AgentConfiguration.get_agent_config_by_key_and_agent_id(session, 'knowledge', agent_id) + agent_execution_knowledge = AgentConfiguration.get_agent_config_by_key_and_agent_id(session= session, key= 'knowledge', agent_id= agent_id) if agent_execution_knowledge: knowledge_name = Knowledges.get_knowledge_from_id(session, int(agent_execution_knowledge.value)).name if knowledge_name is not None: diff --git a/superagi/models/agent_config.py b/superagi/models/agent_config.py index 50d8d195b..94d06c729 100644 --- a/superagi/models/agent_config.py +++ b/superagi/models/agent_config.py @@ -115,12 +115,12 @@ def get_model_api_key(cls, session, agent_id: int, model: str): # if selected_model_source == ModelSourceType.Replicate: # return get_config("REPLICATE_API_TOKEN") # return get_config("OPENAI_API_KEY") -# + @classmethod def get_agent_config_by_key_and_agent_id(cls, session, key: str, agent_id: int): - knowledge_config = session.query(AgentConfiguration).filter( + agent_config = session.query(AgentConfiguration).filter( AgentConfiguration.agent_id == agent_id, AgentConfiguration.key == key ).first() - return knowledge_config \ No newline at end of file + return agent_config \ No newline at end of file diff --git a/tests/unit_tests/apm/test_knowledge_handler.py b/tests/unit_tests/apm/test_knowledge_handler.py index 12ddcca27..27a9bd005 100644 --- a/tests/unit_tests/apm/test_knowledge_handler.py +++ b/tests/unit_tests/apm/test_knowledge_handler.py @@ -17,25 +17,31 @@ def mock_session(): def knowledge_handler(mock_session, organisation_id): return KnowledgeHandler(mock_session, organisation_id) -def test_get_knowledge_wise_usage(knowledge_handler, mock_session): +def test_get_knowledge_usage_by_name(knowledge_handler, mock_session): knowledge_handler.session = mock_session + knowledge_name = 'Knowledge1' mock_knowledge_event = MagicMock() - mock_knowledge_event.knowledge_name = 'Knowledge1' mock_knowledge_event.knowledge_unique_agents = 5 + mock_knowledge_event.knowledge_name = knowledge_name + mock_knowledge_event.id = 1 - mock_session.query.return_value.filter.return_value.group_by.return_value.all.return_value = [mock_knowledge_event] + mock_session.query.return_value.filter_by.return_value.filter.return_value.first.return_value = mock_knowledge_event + mock_session.query.return_value.filter.return_value.group_by.return_value.first.return_value = mock_knowledge_event mock_session.query.return_value.filter.return_value.count.return_value = 10 - result = knowledge_handler.get_knowledge_wise_usage() + result = knowledge_handler.get_knowledge_usage_by_name(knowledge_name) assert isinstance(result, dict) assert result == { - 'Knowledge1': { - 'knowledge_unique_agents': 5, - 'knowledge_calls': 10 - } + 'knowledge_unique_agents': 5, + 'knowledge_calls': 10 } + mock_session.query.return_value.filter_by.return_value.filter.return_value.first.return_value = None + + with pytest.raises(HTTPException): + knowledge_handler.get_knowledge_usage_by_name('NonexistentKnowledge') + def test_get_knowledge_events_by_name(knowledge_handler, mock_session): knowledge_name = 'knowledge1' knowledge_handler.session = mock_session @@ -85,6 +91,7 @@ def test_get_knowledge_events_by_name_knowledge_not_found(knowledge_handler, moc not_found_message = 'Knowledge not found' mock_session.query().filter_by().filter().first.return_value = None + try: knowledge_handler.get_knowledge_events_by_name(knowledge_name) assert False, "Expected HTTPException has not been raised" diff --git a/tests/unit_tests/apm/test_tools_handler.py b/tests/unit_tests/apm/test_tools_handler.py index 4df85a09e..8aac41a94 100644 --- a/tests/unit_tests/apm/test_tools_handler.py +++ b/tests/unit_tests/apm/test_tools_handler.py @@ -64,24 +64,35 @@ def test_get_tool_and_toolkit(tools_handler, mock_session): assert isinstance(output, dict) assert output == {'tool 1': 'toolkit 1'} -def test_get_tool_wise_usage(tools_handler, mock_session): +def test_get_tool_usage_by_name(tools_handler, mock_session): tools_handler.session = mock_session + tool_name = 'Tool1' + formatted_tool_name = tool_name.lower().replace(" ", "") + + mock_tool = MagicMock() + mock_tool.name = tool_name + mock_tool_event = MagicMock() - mock_tool_event.tool_name = 'Tool1' + mock_tool_event.tool_name = formatted_tool_name mock_tool_event.tool_calls = 10 mock_tool_event.tool_unique_agents = 5 + + mock_session.query.return_value.filter_by.return_value.first.return_value = mock_tool + mock_session.query.return_value.filter.return_value.group_by.return_value.first.return_value = mock_tool_event - mock_session.query.return_value.filter.return_value.group_by.return_value = [mock_tool_event] - result = tools_handler.get_tool_wise_usage() - + result = tools_handler.get_tool_usage_by_name(tool_name=tool_name) + assert isinstance(result, dict) assert result == { - 'Tool1': { - 'tool_calls': 10, - 'tool_unique_agents': 5 - } + 'tool_calls': 10, + 'tool_unique_agents': 5 } + mock_session.query.return_value.filter_by.return_value.first.return_value = None + + with pytest.raises(HTTPException): + tools_handler.get_tool_usage_by_name(tool_name="NonexistentTool") + def test_get_tool_events_by_name(tools_handler, mock_session): tool_name = 'tool1' From d6045606dc8ccff8570293f6f4e9dbee592fc767 Mon Sep 17 00:00:00 2001 From: sayan1101 Date: Wed, 30 Aug 2023 10:42:37 +0530 Subject: [PATCH 15/15] toolhandler --- superagi/apm/tools_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superagi/apm/tools_handler.py b/superagi/apm/tools_handler.py index 11e21c882..cc959780e 100644 --- a/superagi/apm/tools_handler.py +++ b/superagi/apm/tools_handler.py @@ -135,7 +135,7 @@ def get_tool_events_by_name(self, tool_name: str) -> List[Dict[str, Union[str, i other_tools = self.session.query( Event.agent_id, - func.array_agg(Event.event_property['tool_name'].astext).label('other_tools') + 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',