From d9d2182f53512d2374d9942047ca32b92bd6d0c8 Mon Sep 17 00:00:00 2001 From: Mustafa Kerem Kurban Date: Fri, 27 Sep 2024 16:44:12 +0200 Subject: [PATCH] add dependencies, and config --- src/neuroagent/app/config.py | 8 +++ src/neuroagent/app/dependencies.py | 18 +++++++ src/neuroagent/app/main.py | 4 +- src/neuroagent/tools/__init__.py | 2 + src/neuroagent/tools/bluenaas_tool.py | 78 +++++++++++++++++++++++++++ 5 files changed, 109 insertions(+), 1 deletion(-) create mode 100644 src/neuroagent/tools/bluenaas_tool.py diff --git a/src/neuroagent/app/config.py b/src/neuroagent/app/config.py index 1e4b01a..fc3f9a4 100644 --- a/src/neuroagent/app/config.py +++ b/src/neuroagent/app/config.py @@ -121,6 +121,13 @@ class SettingsGetMEModel(BaseModel): model_config = ConfigDict(frozen=True) +class SettingsBlueNaaS(BaseModel): + """BlueNaaS settings.""" + + url: str + + model_config = ConfigDict(frozen=True) + class SettingsKnowledgeGraph(BaseModel): """Knowledge graph API settings.""" @@ -166,6 +173,7 @@ class SettingsTools(BaseModel): trace: SettingsTrace = SettingsTrace() kg_morpho_features: SettingsKGMorpho = SettingsKGMorpho() me_model: SettingsGetMEModel = SettingsGetMEModel() + blue_naas: SettingsBlueNaaS = SettingsBlueNaaS() model_config = ConfigDict(frozen=True) diff --git a/src/neuroagent/app/dependencies.py b/src/neuroagent/app/dependencies.py index 430491d..2071792 100644 --- a/src/neuroagent/app/dependencies.py +++ b/src/neuroagent/app/dependencies.py @@ -35,6 +35,7 @@ LiteratureSearchTool, MorphologyFeatureTool, ResolveBrainRegionTool, + BlueNaaSTool, ) from neuroagent.utils import RegionMeta, get_file_from_KG @@ -323,6 +324,23 @@ def get_me_model_tool( ) return tool +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.""" + tool = BlueNaaSTool( + metadata={ + "url": settings.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)], diff --git a/src/neuroagent/app/main.py b/src/neuroagent/app/main.py index 4e331ed..bf5def6 100644 --- a/src/neuroagent/app/main.py +++ b/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}, diff --git a/src/neuroagent/tools/__init__.py b/src/neuroagent/tools/__init__.py index af6761e..ae4ac4e 100644 --- a/src/neuroagent/tools/__init__.py +++ b/src/neuroagent/tools/__init__.py @@ -2,6 +2,7 @@ from neuroagent.tools.electrophys_tool import ElectrophysFeatureTool, FeaturesOutput from neuroagent.tools.get_me_model_tool import GetMEModelTool +from neuroagent.tools.bluenaas_tool import BlueNaaSTool from neuroagent.tools.get_morpho_tool import GetMorphoTool, KnowledgeGraphOutput from neuroagent.tools.kg_morpho_features_tool import ( KGMorphoFeatureOutput, @@ -37,4 +38,5 @@ "ResolveBrainRegionTool", "TracesOutput", "GetMEModelTool", + "BlueNaaSTool", ] diff --git a/src/neuroagent/tools/bluenaas_tool.py b/src/neuroagent/tools/bluenaas_tool.py new file mode 100644 index 0000000..a9ecf82 --- /dev/null +++ b/src/neuroagent/tools/bluenaas_tool.py @@ -0,0 +1,78 @@ +"""BlueNaaS single cell stimulation, simulation and synapse placement tool.""" +from typing import List, Literal, Optional +from pydantic import BaseModel, Field + +import logging +from typing import Any, Type +from pydantic import BaseModel +from neuroagent.tools.base_tool import BaseToolOutput, BasicTool +from neuroagent.utils import get_kg_data + +logger = logging.getLogger(__name__) + +class SynapseConfig(BaseModel): + id: str + delay: int + duration: int + frequency: int + weightScalar: float + +class StimulusConfig(BaseModel): + stimulusType: Literal["current_clamp", "voltage_clamp"] + stimulusProtocol: Literal["ap_waveform", "idrest", "iv", "fire_pattern", "pos_cheops", "neg_cheops"] + amplitudes: List[float] + +class CurrentInjectionConfig(BaseModel): + injectTo: str + stimulus: StimulusConfig + +class RecordFromConfig(BaseModel): + section: str + offset: float + +class ConditionsConfig(BaseModel): + celsius: float + vinit: float + hypamp: float + max_time: int + time_step: float + seed: int + +class InputBlueNaaS(BaseModel): + model_id: str + synapses: List[SynapseConfig] + currentInjection: CurrentInjectionConfig + recordFrom: List[RecordFromConfig] + conditions: ConditionsConfig + type: Literal["single-neuron-simulation"] + simulationDuration: int + +class BlueNaaSOutput(BaseModel): + status: str + result: Optional[dict] + error: Optional[str] + + +class BlueNaaSTool(BasicTool): + name: str = "blue-naas-tool" + description: str = """Runs a single-neuron simulation using the BlueNaaS service. + Requires a 'model_id' which can be fetched using the 'get-me-model-tool'. + The input configuration should be provided by the user otherwise agent will probe the user + with the selected default values.""" + metadata: dict[str, Any] + args_schema: Type[BaseModel] = InputBlueNaaS + + async def _arun(self, **kwargs) -> BlueNaaSOutput: + """Run a single-neuron simulation using the BlueNaaS service.""" + logger.info(f"Running BlueNaaS tool with inputs: {kwargs}") + try: + response = await self.metadata["httpx_client"].post( + url=self.metadata["url"], + headers={"Authorization": f"Bearer {self.metadata['token']}"}, + json=kwargs, + ) + response_data = response.json() + return BlueNaaSOutput(status="success", result=response_data) + except Exception as e: + logger.error(f"Error running BlueNaaS tool: {e}") + return BlueNaaSOutput(status="error", error=str(e)) \ No newline at end of file