From 5ea2c69ccfe62d918d6abe7828e31082daba43ff Mon Sep 17 00:00:00 2001
From: Nicolas Frank <nicolas@frank.nom.fr>
Date: Wed, 8 Jan 2025 14:31:57 +0100
Subject: [PATCH] Fix mypy

---
 src/neuroagent/app/routers/qa.py | 13 +++++++++++--
 src/neuroagent/new_types.py      |  2 +-
 src/neuroagent/stream.py         |  4 ++--
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/neuroagent/app/routers/qa.py b/src/neuroagent/app/routers/qa.py
index 157b446..0efc2dd 100644
--- a/src/neuroagent/app/routers/qa.py
+++ b/src/neuroagent/app/routers/qa.py
@@ -45,7 +45,16 @@ async def run_simple_agent(
 ) -> AgentResponse:
     """Run a single agent query."""
     response = await agent_routine.arun(
-        agent, [{"role": "user", "content": user_request.query}], context_variables
+        agent,
+        [
+            Messages(
+                order=0,
+                thread_id="Dummy_thread_id",
+                entity=Entity.USER,
+                content=json.dumps({"role": "user", "content": user_request.query}),
+            )
+        ],
+        context_variables,
     )
     return AgentResponse(message=response.messages[-1]["content"])
 
@@ -80,7 +89,7 @@ async def run_chat_agent(
     thread.update_date = utc_now()
     await session.commit()
 
-    if response.hil_messages:
+    if response.hil_messages is not None:
         return response.hil_messages
     return AgentResponse(message=response.messages[-1]["content"])
 
diff --git a/src/neuroagent/new_types.py b/src/neuroagent/new_types.py
index ffe9cd6..498b695 100644
--- a/src/neuroagent/new_types.py
+++ b/src/neuroagent/new_types.py
@@ -41,7 +41,7 @@ class Response(BaseModel):
     messages: list[dict[str, Any]] = []
     agent: Agent | None = None
     context_variables: dict[str, Any] = {}
-    hil_messages: list[HILResponse | None] = []
+    hil_messages: list[HILResponse] | None = None
 
 
 class AgentRequest(BaseModel):
diff --git a/src/neuroagent/stream.py b/src/neuroagent/stream.py
index 98a0165..2bbef4e 100644
--- a/src/neuroagent/stream.py
+++ b/src/neuroagent/stream.py
@@ -8,7 +8,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from neuroagent.agent_routine import AgentsRoutine
-from neuroagent.app.database.sql_schemas import Messages, utc_now
+from neuroagent.app.database.sql_schemas import Messages, Threads, utc_now
 from neuroagent.new_types import Agent, Response
 
 
@@ -17,7 +17,7 @@ async def stream_agent_response(
     agent: Agent,
     messages: list[Messages],
     context_variables: dict[str, Any],
-    thread: str,
+    thread: Threads,
     request: Request,
 ) -> AsyncIterator[str]:
     """Redefine fastAPI connections to enable streaming."""