Issues fixing (#1330)

* fix call_log_helper.py and updated test_call_log_helper.py --------- Co-authored-by: Fluder-Paradyne <[email protected]>
TransformerOptimus · Oct 10, 2023 · adb8e41 · adb8e41
1 parent c02fe30
commit adb8e41
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 14 deletions.
diff --git a/superagi/apm/call_log_helper.py b/superagi/apm/call_log_helper.py
@@ -52,22 +52,25 @@ 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()
 
-                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()
+            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}
 
+            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.name if agent is not None else None,
+                    'agent_name': agent_id_name_map[run.agent_id] if run.agent_id in agent_id_name_map else None,
                     'tokens_consumed': run.tokens_consumed,
                     'tool_used': run.tool_used,
-                    'toolkit_name': toolkit.name if toolkit is not None else None,
+                    'toolkit_name': tools_name_toolkit_id_map[run.tool_used] if run.tool_used in tools_name_toolkit_id_map else None,
                     'org_id': run.org_id,
                     'created_at': run.created_at,
                     'updated_at': run.updated_at,
@@ -79,4 +82,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,13 +55,12 @@ def test_fetch_data_success(call_log_helper, mock_session):
         model='test_model',
         org_id=1
     )]
-    agent = Agent(name='test_agent')
-    tool = Tool(name='test_tool', toolkit_id=1)
-    toolkit = Toolkit(name='test_toolkit')
+    agents = [Agent(name='test_agent')]
+    tools = [Tool(name='test_tool', toolkit_id=1)]
+    toolkits = [Toolkit(name='test_toolkit')]
 
     # setup return values for the mock methods
-    mock_session.query().filter().first.side_effect = [summary_result, agent, tool, toolkit]
-    mock_session.query().filter().all.return_value = runs
+    mock_session.query().filter().first.side_effect = [summary_result, runs, agents, toolkits, tools]
 
     result = call_log_helper.fetch_data('test_model')