Revert "Issues fixing (#1330)"

This reverts commit adb8e41.
TransformerOptimus · Oct 10, 2023 · 71e052b · 71e052b
1 parent adb8e41
commit 71e052b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 16 deletions.
diff --git a/superagi/apm/call_log_helper.py b/superagi/apm/call_log_helper.py
@@ -52,25 +52,22 @@ def fetch_data(self, model: str):
 
             runs = self.session.query(CallLogs).filter(CallLogs.model == model,
                                                        CallLogs.org_id == self.organisation_id).all()
+            for run in runs:
+                agent = self.session.query(Agent).filter(Agent.id == run.agent_id).first()
 
-            run_agent_ids = [run.agent_id for run in runs]
-            agents = self.session.query(Agent).filter(Agent.id.in_(run_agent_ids)).all()
-            agent_id_name_map = {agent.id: agent.name for agent in agents}
-            tools_used = [run.tool_used for run in runs]
-            toolkit_ids_allowed = self.session.query(Toolkit.id).filter(Toolkit.organisation_id == self.organisation_id).all()
-            tools = self.session.query(Tool).filter(Tool.id.in_(tools_used), Tool.toolkit_id.in_(toolkit_ids_allowed))\
-                .all()
-            tools_name_toolkit_id_map = {tool.name: tool.toolkit_id for tool in tools}
+                toolkit = None
+                tool = self.session.query(Tool).filter(Tool.name == run.tool_used).first()
+                if tool:
+                    toolkit = self.session.query(Toolkit).filter(Toolkit.id == tool.toolkit_id).first()
 
-            for run in runs:
                 model_data['runs'].append({
                     'id': run.id,
                     'agent_execution_name': run.agent_execution_name,
                     'agent_id': run.agent_id,
-                    'agent_name': agent_id_name_map[run.agent_id] if run.agent_id in agent_id_name_map else None,
+                    'agent_name': agent.name if agent is not None else None,
                     'tokens_consumed': run.tokens_consumed,
                     'tool_used': run.tool_used,
-                    'toolkit_name': tools_name_toolkit_id_map[run.tool_used] if run.tool_used in tools_name_toolkit_id_map else None,
+                    'toolkit_name': toolkit.name if toolkit is not None else None,
                     'org_id': run.org_id,
                     'created_at': run.created_at,
                     'updated_at': run.updated_at,
@@ -82,4 +79,4 @@ def fetch_data(self, model: str):
 
         except SQLAlchemyError as err:
             logging.error(f"Error while fetching call log data: {str(err)}")
-            return None
+            return None
diff --git a/tests/unit_tests/apm/test_call_log_helper.py b/tests/unit_tests/apm/test_call_log_helper.py
@@ -55,12 +55,13 @@ def test_fetch_data_success(call_log_helper, mock_session):
         model='test_model',
         org_id=1
     )]
-    agents = [Agent(name='test_agent')]
-    tools = [Tool(name='test_tool', toolkit_id=1)]
-    toolkits = [Toolkit(name='test_toolkit')]
+    agent = Agent(name='test_agent')
+    tool = Tool(name='test_tool', toolkit_id=1)
+    toolkit = Toolkit(name='test_toolkit')
 
     # setup return values for the mock methods
-    mock_session.query().filter().first.side_effect = [summary_result, runs, agents, toolkits, tools]
+    mock_session.query().filter().first.side_effect = [summary_result, agent, tool, toolkit]
+    mock_session.query().filter().all.return_value = runs
 
     result = call_log_helper.fetch_data('test_model')