diff --git a/agents-api/agents_api/models/entry/proc_mem_context.py b/agents-api/agents_api/models/entry/proc_mem_context.py
index 82b114e75..4b1d650d7 100644
--- a/agents-api/agents_api/models/entry/proc_mem_context.py
+++ b/agents-api/agents_api/models/entry/proc_mem_context.py
@@ -183,7 +183,7 @@ def proc_mem_context_query(
                 created_at,
             }},
             ~information_snippets:embedding_space {{
-                doc_id,
+                doc_id: agent_doc_id,
                 snippet_idx,
                 title,
                 snippet |
@@ -210,7 +210,7 @@ def proc_mem_context_query(
                 created_at,
             }},
             ~information_snippets:embedding_space {{
-                doc_id,
+                doc_id: user_doc_id,
                 snippet_idx,
                 title,
                 snippet |