type(data profile): add llm_call_manager for each mode

Author: StCarmen

alias/src/alias/agent/agents/data_source/_data_profiler_factory.py

@@ -97,7 +97,7 @@ async def generate_profile(self) -> Dict[str, Any]:
                 self.data,
             )
             # content = self.prompt.format(data=self.data)
-            res = await self._call_model(content)
+            res = await self._generate_profile_by_llm(content)
             self.profile = self._wrap_data_response(res)
         except Exception as e:
             logger.warning(f"Error generating profile: {e}")
@@ -157,7 +157,7 @@ async def _read_data(self):
             Processed data in appropriate format for the data type
         """
 
-    async def _call_model(
+    async def _generate_profile_by_llm(
         self,
         content: Any,
     ) -> Dict[str, Any]:
@@ -349,7 +349,7 @@ async def _extract_irregular_table(
         """
         prompt = self._load_prompt("IRREGULAR")
         content = prompt.format(raw_snippet_data=raw_data_snippet)
-        res = await self._call_model(content=content)
+        res = await self._generate_profile_by_llm(content=content)
 
         if "is_extractable_table" in res and res["is_extractable_table"]:
             logger.debug(res["reasoning"])

alias/src/alias/agent/run.py

@@ -248,7 +248,18 @@ async def arun_deepresearch_agent(
         "run_shell_command",
     ]
     share_tools(global_toolkit, worker_toolkit, test_tool_list)
-    await prepare_data_sources(session_service, sandbox, worker_toolkit)
+
+    llm_call_manager = LLMCallManager(
+        base_model_name=MODEL_CONFIG_NAME,
+        vl_model_name=VL_MODEL_NAME,
+        model_formatter_mapping=MODEL_FORMATTER_MAPPING,
+    )
+    await prepare_data_sources(
+        session_service,
+        sandbox,
+        worker_toolkit,
+        llm_call_manager,
+    )
 
     worker_agent = DeepResearchAgent(
         name="Deep_Research_Agent",
@@ -315,7 +326,18 @@ async def arun_finance_agent(
         description="Finance Analysis tools",
         active=True,
     )
-    await prepare_data_sources(session_service, sandbox, worker_toolkit)
+
+    llm_call_manager = LLMCallManager(
+        base_model_name=MODEL_CONFIG_NAME,
+        vl_model_name=VL_MODEL_NAME,
+        model_formatter_mapping=MODEL_FORMATTER_MAPPING,
+    )
+    await prepare_data_sources(
+        session_service,
+        sandbox,
+        worker_toolkit,
+        llm_call_manager,
+    )
 
     worker_agent = DeepResearchAgent(
         name="Deep_Research_Agent",
@@ -425,7 +447,17 @@ async def arun_browseruse_agent(
         add_all=True,
         is_browser_toolkit=True,
     )
-    await prepare_data_sources(session_service, sandbox, browser_toolkit)
+    llm_call_manager = LLMCallManager(
+        base_model_name=MODEL_CONFIG_NAME,
+        vl_model_name=VL_MODEL_NAME,
+        model_formatter_mapping=MODEL_FORMATTER_MAPPING,
+    )
+    await prepare_data_sources(
+        session_service,
+        sandbox,
+        browser_toolkit,
+        llm_call_manager,
+    )
 
     logger.info("Init browser toolkit")
     try: