Bluenaas tool

.env.example

@@ -23,6 +23,8 @@ NEUROAGENT_TOOLS__LITERATURE__USE_RERANKER=
 
 NEUROAGENT_TOOLS__KNOWLEDGE_GRAPH__SEARCH_SIZE=
 
+NEUROAGENT_TOOLS__BLUENAAS__URL=
+
 NEUROAGENT_TOOLS__TRACE__SEARCH_SIZE=
 
 NEUROAGENT_TOOLS__KG_MORPHO__SEARCH_SIZE=

.github/workflows/ci.yaml

@@ -34,7 +34,7 @@ jobs:
       - name: Set up environment
         run: |
           pip install --upgrade pip wheel setuptools
-          pip install bandit[toml]==1.7.4 ruff==0.5.5
+          pip install bandit[toml]==1.7.4 ruff==0.6.7
       - name: Linting check
         run: |
           bandit -qr -c pyproject.toml src/

.gitignore

@@ -6,7 +6,7 @@
 .python-version
 .vscode
 .pypirc
-
+src/neuroagent/app/data
 # Version file
 
 # Byte-compiled / optimized / DLL files
@@ -152,3 +152,8 @@ cython_debug/
 # static files generated from Django application using `collectstatic`
 media
 static
+
+# database stuff
+*db
+*.db-shm
+*.db-wal

CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Add get morphoelectric (me) model tool
+- Filter me model by username
+- Add bluenaas tool
+- Add hierarcichal multi agent for bluenass
+
 ## [0.1.1] - 26.09.2024
 
 ### Fixed

src/neuroagent/agents/__init__.py

@@ -3,6 +3,7 @@
 from neuroagent.agents.base_agent import AgentOutput, AgentStep, BaseAgent
 from neuroagent.agents.simple_agent import SimpleAgent
 from neuroagent.agents.simple_chat_agent import SimpleChatAgent
+# from neuroagent.agents.bluenaas_sim_agent import BluenaasSimAgent
 
 __all__ = [
     "AgentOutput",

src/neuroagent/app/config.py

@@ -6,18 +6,17 @@
 
 from dotenv import dotenv_values
 from fastapi.openapi.models import OAuthFlowPassword, OAuthFlows
-from pydantic import BaseModel, ConfigDict, SecretStr, model_validator
+from pydantic import BaseModel, ConfigDict, SecretStr, model_validator, Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 
 class SettingsAgent(BaseModel):
     """Agent setting."""
 
-    model: Literal["simple", "multi"] = "simple"
+    model: Literal["simple", "multi-supervisor","multi-hierarchical"] = "simple"
 
     model_config = ConfigDict(frozen=True)
 
-
 class SettingsDB(BaseModel):
     """DB settings for retrieving history."""
 
@@ -114,6 +113,20 @@ class SettingsGetMorpho(BaseModel):
     model_config = ConfigDict(frozen=True)
 
 
+class SettingsGetMEModel(BaseModel):
+    """Get ME Model settings."""
+
+    search_size: int = 10
+
+    model_config = ConfigDict(frozen=True)
+
+
+class SettingsBlueNaaS(BaseModel):
+    """BlueNaaS settings."""
+
+    url: str = "https://openbluebrain.com/api/bluenaas/single-neuron/run" #TODO: move to .env
+    model_config = ConfigDict(frozen=True)
+
 class SettingsKnowledgeGraph(BaseModel):
     """Knowledge graph API settings."""
 
@@ -157,6 +170,8 @@ class SettingsTools(BaseModel):
     morpho: SettingsGetMorpho = SettingsGetMorpho()
     trace: SettingsTrace = SettingsTrace()
     kg_morpho_features: SettingsKGMorpho = SettingsKGMorpho()
+    me_model: SettingsGetMEModel = SettingsGetMEModel()
+    bluenaas: SettingsBlueNaaS = SettingsBlueNaaS()
 
     model_config = ConfigDict(frozen=True)
 
@@ -206,6 +221,7 @@ class Settings(BaseSettings):
     keycloak: SettingsKeycloak = SettingsKeycloak()  # has no required
     misc: SettingsMisc = SettingsMisc()  # has no required
 
+
     model_config = SettingsConfigDict(
         env_file=".env",
         env_prefix="NEUROAGENT_",

src/neuroagent/app/dependencies.py

@@ -25,15 +25,17 @@
 from neuroagent.agents.base_agent import AsyncSqliteSaverWithPrefix
 from neuroagent.app.config import Settings
 from neuroagent.cell_types import CellTypesMeta
-from neuroagent.multi_agents import BaseMultiAgent, SupervisorMultiAgent
+from neuroagent.multi_agents import BaseMultiAgent, SupervisorMultiAgent, HierarchicalTeamAgent
 from neuroagent.tools import (
     ElectrophysFeatureTool,
+    GetMEModelTool,
     GetMorphoTool,
     GetTracesTool,
     KGMorphoFeatureTool,
     LiteratureSearchTool,
     MorphologyFeatureTool,
     ResolveBrainRegionTool,
+    BlueNaaSTool,
 )
 from neuroagent.utils import RegionMeta, get_file_from_KG
 
@@ -304,6 +306,44 @@ def get_morphology_feature_tool(
     return tool
 
 
+def get_me_model_tool(
+    settings: Annotated[Settings, Depends(get_settings)],
+    token: Annotated[str, Depends(get_kg_token)],
+    httpx_client: Annotated[AsyncClient, Depends(get_httpx_client)],
+) -> GetMEModelTool:
+    """Load get ME model tool."""
+    tool = GetMEModelTool(
+        metadata={
+            "url": settings.knowledge_graph.url,
+            "token": token,
+            "httpx_client": httpx_client,
+            "search_size": settings.tools.me_model.search_size,
+            "brainregion_path": settings.knowledge_graph.br_saving_path,
+            "celltypes_path": settings.knowledge_graph.ct_saving_path,
+        }
+    )
+    return tool
+
+async def run_single_cell_sim_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."""
+    # Run GetMEModelTool to fetch me_model_id
+    # Create BlueNaaSTool with the fetched me_model_id
+    tool = BlueNaaSTool(
+        metadata={
+            "url": settings.tools.bluenaas.url,
+            "token": token,
+            "httpx_client": httpx_client,
+        }
+    )
+    return tool
+
+def get_project_id():
+    pass
+
 def get_language_model(
     settings: Annotated[Settings, Depends(get_settings)],
 ) -> ChatOpenAI:
@@ -315,7 +355,7 @@ def get_language_model(
         openai_api_key=settings.openai.token.get_secret_value(),  # type: ignore
         max_tokens=settings.openai.max_tokens,
         seed=78,
-        streaming=True,
+        streaming=False,
     )
 
 
@@ -369,11 +409,15 @@ def get_agent(
         ElectrophysFeatureTool, Depends(get_electrophys_feature_tool)
     ],
     traces_tool: Annotated[GetTracesTool, Depends(get_traces_tool)],
+    me_model_tool: Annotated[GetMEModelTool, Depends(get_me_model_tool)],
+    bluenaas_tool: Annotated[BlueNaaSTool, Depends(run_single_cell_sim_tool)],
+    # bluenaas_sim_agent: Annotated[BluenaasSimAgent, Depends(get_bluenaas_sim_agent)],  # Include BluenaasSimAgent
+
     settings: Annotated[Settings, Depends(get_settings)],
 ) -> BaseAgent | BaseMultiAgent:
     """Get the generative question answering service."""
-    if settings.agent.model == "multi":
-        logger.info("Load multi-agent chat")
+    if settings.agent.model == "multi-supervisor":
+        logger.info("Load multi-agent (flat) chat")
         tools_list = [
             ("literature", [literature_tool]),
             (
@@ -383,11 +427,32 @@ def get_agent(
                     morpho_tool,
                     morphology_feature_tool,
                     kg_morpho_feature_tool,
+                    literature_tool,
+                    electrophys_feature_tool,
+                    traces_tool,
+                    me_model_tool,
+                    bluenaas_tool,
                 ],
             ),
             ("traces", [br_resolver_tool, electrophys_feature_tool, traces_tool]),
         ]
         return SupervisorMultiAgent(llm=llm, agents=tools_list)  # type: ignore
+
+    elif settings.agent.model == "multi-hierarchical":
+        logger.info("Load multi-agent (hierarchical teams) chat")
+        tools_list = [                 
+                    br_resolver_tool,
+                    morpho_tool,
+                    morphology_feature_tool,
+                    kg_morpho_feature_tool,
+                    literature_tool,
+                    electrophys_feature_tool,
+                    traces_tool,
+                    me_model_tool,
+                    bluenaas_tool,
+                      ]
+        return HierarchicalTeamAgent(llm=llm, agents=tools_list)  # type: ignore
+
     else:
         tools = [
             literature_tool,
@@ -397,6 +462,9 @@ def get_agent(
             kg_morpho_feature_tool,
             electrophys_feature_tool,
             traces_tool,
+            me_model_tool,
+            bluenaas_tool,
+            # bluenaas_sim_agent,
         ]
         logger.info("Load simple agent")
         return SimpleAgent(llm=llm, tools=tools)  # type: ignore
@@ -420,20 +488,45 @@ def get_chat_agent(
         ElectrophysFeatureTool, Depends(get_electrophys_feature_tool)
     ],
     traces_tool: Annotated[GetTracesTool, Depends(get_traces_tool)],
+    me_model_tool: Annotated[GetMEModelTool, Depends(get_me_model_tool)],
+    bluenaas_tool: Annotated[BlueNaaSTool, Depends(run_single_cell_sim_tool)],
+    # bluenaas_sim_agent: Annotated[BluenaasSimAgent, Depends(get_bluenaas_sim_agent)], 
+    settings: Annotated[Settings, Depends(get_settings)],
+
 ) -> BaseAgent:
     """Get the generative question answering service."""
-    logger.info("Load simple chat")
-    tools = [
-        literature_tool,
-        br_resolver_tool,
-        morpho_tool,
-        morphology_feature_tool,
-        kg_morpho_feature_tool,
-        electrophys_feature_tool,
-        traces_tool,
-    ]
-    return SimpleChatAgent(llm=llm, tools=tools, memory=memory)  # type: ignore
 
+    if settings.agent.model == "multi-hierarchical":
+        logger.info("Load multi-agent (hierarchical teams) chat")
+        tools = {               
+            "br_resolver_tool": br_resolver_tool,
+            "morpho_tool": morpho_tool,
+            "morphology_feature_tool": morphology_feature_tool,
+            "kg_morpho_feature_tool": kg_morpho_feature_tool,
+            "literature_tool": literature_tool,
+            "electrophys_feature_tool": electrophys_feature_tool,
+            "traces_tool": traces_tool,
+            "me_model_tool": me_model_tool,
+            "bluenaas_tool": bluenaas_tool,
+        }
+        return HierarchicalTeamAgent(llm=llm, tools=tools, memory=memory)  # type: ignore
+    elif settings.agent.model == "simple":
+        logger.info("Load simple chat")
+        tools = [
+            literature_tool,
+            br_resolver_tool,
+            morpho_tool,
+            morphology_feature_tool,
+            kg_morpho_feature_tool,
+            electrophys_feature_tool,
+            traces_tool,
+            me_model_tool,
+            bluenaas_tool,
+        ]
+        return SimpleChatAgent(llm=llm, tools=tools, memory=memory)  # type: ignore
+
+    else:
+        logger.error("Invalid agent model for chat agent.")
 
 async def get_update_kg_hierarchy(
     token: Annotated[str, Depends(get_kg_token)],

src/neuroagent/app/main.py

@@ -23,6 +23,7 @@
     get_kg_token,
     get_settings,
     get_update_kg_hierarchy,
+    get_project_id,
 )
 from neuroagent.app.middleware import strip_path_prefix
 from neuroagent.app.routers import qa
@@ -125,7 +126,8 @@ async def lifespan(fastapi_app: FastAPI) -> AsyncContextManager[None]:  # type:
     title="Agents",
     summary=(
         "Use an AI agent to answer queries based on the knowledge graph, literature"
-        " search and neuroM."
+        " search and neuroM. Extract statistical information from morphologies and electrophysiology."
+        "Run single cell simulations with BlueNaaS API."
     ),
     version=__version__,
     swagger_ui_parameters={"tryItOutEnabled": True},

src/neuroagent/multi_agents/__init__.py

@@ -2,8 +2,9 @@
 
 from neuroagent.multi_agents.base_multi_agent import BaseMultiAgent
 from neuroagent.multi_agents.supervisor_multi_agent import SupervisorMultiAgent
-
+from neuroagent.multi_agents.hierarchical_multi_agent import HierarchicalTeamAgent
 __all__ = [
     "BaseMultiAgent",
     "SupervisorMultiAgent",
+    "HierarchicalTeamAgent",
 ]

src/neuroagent/multi_agents/base_multi_agent.py

@@ -14,8 +14,8 @@ class BaseMultiAgent(BaseModel, ABC):
     """Base class for multi agents."""
 
     llm: BaseChatModel
-    main_agent: Any
-    agents: list[tuple[str, list[BasicTool]]]
+    main_agent: Any | None = None 
+    agents: list[tuple[str, list[BasicTool]]] | None = None
 
     model_config = ConfigDict(arbitrary_types_allowed=True)