Fix mypy

BlueBrain · Sep 24, 2024 · 90594e9 · 90594e9
1 parent 1cc4751
commit 90594e9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/neuroagent/agents/simple_chat_agent.py b/src/neuroagent/agents/simple_chat_agent.py
@@ -17,7 +17,7 @@
 class SimpleChatAgent(BaseAgent):
     """Simple Agent class."""
 
-    memory: BaseCheckpointSaver
+    memory: BaseCheckpointSaver[Any]
 
     @model_validator(mode="before")
     @classmethod

diff --git a/src/neuroagent/app/dependencies.py b/src/neuroagent/app/dependencies.py
@@ -321,7 +321,7 @@ def get_language_model(
 
 async def get_agent_memory(
     connection_string: Annotated[str | None, Depends(get_connection_string)],
-) -> AsyncIterator[BaseCheckpointSaver | None]:
+) -> AsyncIterator[BaseCheckpointSaver[Any] | None]:
     """Get the agent checkpointer."""
     if connection_string:
         if connection_string.startswith("sqlite"):
@@ -404,7 +404,7 @@ def get_agent(
 
 def get_chat_agent(
     llm: Annotated[ChatOpenAI, Depends(get_language_model)],
-    memory: Annotated[BaseCheckpointSaver, Depends(get_agent_memory)],
+    memory: Annotated[BaseCheckpointSaver[Any], Depends(get_agent_memory)],
     literature_tool: Annotated[LiteratureSearchTool, Depends(get_literature_tool)],
     br_resolver_tool: Annotated[
         ResolveBrainRegionTool, Depends(get_brain_region_resolver_tool)