diff --git a/agents-api/agents_api/queries/chat/prepare_chat_context.py b/agents-api/agents_api/queries/chat/prepare_chat_context.py
index 7b8255633..d88c66a83 100644
--- a/agents-api/agents_api/queries/chat/prepare_chat_context.py
+++ b/agents-api/agents_api/queries/chat/prepare_chat_context.py
@@ -7,6 +7,7 @@
 from ..session.prepare_session_data import prepare_session_data
 from ..utils import (
     wrap_in_class,
+    pg_query,
 )
 
 ModelT = TypeVar("ModelT", bound=Any)
@@ -64,7 +65,7 @@
         ],
     },
 )
-@cozo_query
+@pg_query
 @beartype
 def prepare_chat_context(
     *,