mypy

src/neuroagent/agents/base_agent.py

@@ -93,7 +93,7 @@ async def arun(self, *args: Any, **kwargs: Any) -> AgentOutput:
         """Arun method of the service."""
 
     @abstractmethod
-    async def astream(self, *args: Any, **kwargs: Any) -> AsyncIterator[str]:
+    def astream(self, *args: Any, **kwargs: Any) -> AsyncIterator[str]:
         """Astream method of the service."""
 
     @staticmethod

src/neuroagent/agents/simple_agent.py

@@ -58,7 +58,7 @@ async def arun(self, query: str) -> Any:
         result = await self.agent.ainvoke({"messages": [("human", query)]})
         return self._process_output(result)
 
-    async def astream(self, query: str) -> AsyncIterator[str]:  # type: ignore
+    async def astream(self, query: str) -> AsyncIterator[str]:
         """Run the agent against a query in streaming way.
 
         Parameters

src/neuroagent/agents/simple_chat_agent.py

@@ -46,7 +46,7 @@ async def arun(self, thread_id: str, query: str) -> Any:
 
     async def astream(
         self, thread_id: str, query: str, connection_string: str | None = None
-    ) -> AsyncIterator[str]:  # type: ignore
+    ) -> AsyncIterator[str]:
         """Run the agent against a query in streaming way.
 
         Parameters
@@ -58,8 +58,8 @@ async def astream(
         connection_string
             connection string for the checkpoint database.
 
-        Returns
-        -------
+        Yields
+        ------
             Iterator streaming the processed output of the LLM
         """
         async with (

src/neuroagent/app/routers/qa.py

@@ -69,4 +69,4 @@ async def run_streamed_chat_agent(
             thread_id=thread_id,
             connection_string=connection_string,
         )
-    )  # type: ignore
+    )