Add tool to dependencies

BlueBrain · Oct 7, 2024 · 9862ea6 · 9862ea6
1 parent fad474f
commit 9862ea6
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/src/neuroagent/app/dependencies.py b/src/neuroagent/app/dependencies.py
@@ -27,6 +27,7 @@
 from neuroagent.cell_types import CellTypesMeta
 from neuroagent.multi_agents import BaseMultiAgent, SupervisorMultiAgent
 from neuroagent.tools import (
+    BlueNaaSTool,
     ElectrophysFeatureTool,
     GetMEModelTool,
     GetMorphoTool,
@@ -182,6 +183,22 @@ def get_kg_token(
             return settings.knowledge_graph.token.get_secret_value()  # type: ignore
 
 
+def get_bluenaas_tool(
+    settings: Annotated[Settings, Depends(get_settings)],
+    token: Annotated[str, Depends(get_kg_token)],
+    httpx_client: Annotated[AsyncClient, Depends(get_httpx_client)],
+) -> BlueNaaSTool:
+    """Load BlueNaaS tool."""
+    tool = BlueNaaSTool(
+        metadata={
+            "url": settings.tools.bluenaas.url,
+            "token": token,
+            "httpx_client": httpx_client,
+        }
+    )
+    return tool
+
+
 def get_literature_tool(
     token: Annotated[str, Depends(get_kg_token)],
     settings: Annotated[Settings, Depends(get_settings)],
@@ -374,6 +391,7 @@ async def get_agent_memory(
 
 def get_agent(
     llm: Annotated[ChatOpenAI, Depends(get_language_model)],
+    blue_naas_tool: Annotated[BlueNaaSTool, Depends(get_bluenaas_tool)],
     literature_tool: Annotated[LiteratureSearchTool, Depends(get_literature_tool)],
     br_resolver_tool: Annotated[
         ResolveBrainRegionTool, Depends(get_brain_region_resolver_tool)
@@ -411,6 +429,7 @@ def get_agent(
         return SupervisorMultiAgent(llm=llm, agents=tools_list)  # type: ignore
     else:
         tools = [
+            blue_naas_tool,
             literature_tool,
             br_resolver_tool,
             morpho_tool,
@@ -427,6 +446,7 @@ def get_agent(
 def get_chat_agent(
     llm: Annotated[ChatOpenAI, Depends(get_language_model)],
     memory: Annotated[BaseCheckpointSaver[Any], Depends(get_agent_memory)],
+    blue_naas_tool: Annotated[BlueNaaSTool, Depends(get_bluenaas_tool)],
     literature_tool: Annotated[LiteratureSearchTool, Depends(get_literature_tool)],
     br_resolver_tool: Annotated[
         ResolveBrainRegionTool, Depends(get_brain_region_resolver_tool)
@@ -446,6 +466,7 @@ def get_chat_agent(
     """Get the generative question answering service."""
     logger.info("Load simple chat")
     tools = [
+        blue_naas_tool,
         literature_tool,
         br_resolver_tool,
         morpho_tool,

diff --git a/src/neuroagent/tools/__init__.py b/src/neuroagent/tools/__init__.py
@@ -1,5 +1,6 @@
 """Tools folder."""
 
+from neuroagent.tools.bluenaas_tool import BlueNaaSTool
 from neuroagent.tools.electrophys_tool import ElectrophysFeatureTool, FeaturesOutput
 from neuroagent.tools.get_me_model_tool import GetMEModelTool
 from neuroagent.tools.get_morpho_tool import GetMorphoTool, KnowledgeGraphOutput
@@ -22,6 +23,7 @@
 from neuroagent.tools.traces_tool import GetTracesTool, TracesOutput
 
 __all__ = [
+    "BlueNaaSTool",
     "BRResolveOutput",
     "ElectrophysFeatureTool",
     "FeaturesOutput",