From 70e5781c3b4c35a9702ebd894de2d985670e1cdd Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Tue, 21 Apr 2026 12:33:36 +0100
Subject: [PATCH] Fix RAG

---
 .../trustgraph/retrieval/document_rag/document_rag.py  | 10 ++++++----
 .../trustgraph/retrieval/document_rag/rag.py           |  2 ++
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py b/trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py
index dfe4e051..283cf013 100644
--- a/trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py
+++ b/trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py
@@ -26,10 +26,11 @@ LABEL="http://www.w3.org/2000/01/rdf-schema#label"
 class Query:
 
     def __init__(
-            self, rag, user, collection, verbose,
+            self, rag, workspace, user, collection, verbose,
             doc_limit=20, track_usage=None,
     ):
         self.rag = rag
+        self.workspace = workspace
         self.user = user
         self.collection = collection
         self.verbose = verbose
@@ -122,7 +123,7 @@ class Query:
         for match in chunk_matches:
             if match.chunk_id:
                 try:
-                    content = await self.rag.fetch_chunk(match.chunk_id, self.user)
+                    content = await self.rag.fetch_chunk(match.chunk_id, self.workspace)
                     docs.append(content)
                     chunk_ids.append(match.chunk_id)
                 except Exception as e:
@@ -154,7 +155,7 @@ class DocumentRag:
             logger.debug("DocumentRag initialized")
 
     async def query(
-            self, query, user="trustgraph", collection="default",
+            self, query, workspace="default", user="", collection="default",
             doc_limit=20, streaming=False, chunk_callback=None,
             explain_callback=None, save_answer_callback=None,
     ):
@@ -210,7 +211,8 @@ class DocumentRag:
             await explain_callback(q_triples, q_uri)
 
         q = Query(
-            rag=self, user=user, collection=collection, verbose=self.verbose,
+            rag=self, workspace=workspace, user=user, collection=collection,
+            verbose=self.verbose,
             doc_limit=doc_limit, track_usage=track_usage,
         )
 
diff --git a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py
index 6dae21fd..f43e96d5 100755
--- a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py
+++ b/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py
@@ -201,6 +201,7 @@ class Processor(FlowProcessor):
                 # All chunks (including final one with end_of_stream=True) are sent via callback
                 response, usage = await self.rag.query(
                     v.query,
+                    workspace=flow.workspace,
                     user=v.user,
                     collection=v.collection,
                     doc_limit=doc_limit,
@@ -226,6 +227,7 @@ class Processor(FlowProcessor):
                 # Non-streaming path - single response with answer and token usage
                 response, usage = await self.rag.query(
                     v.query,
+                    workspace=flow.workspace,
                     user=v.user,
                     collection=v.collection,
                     doc_limit=doc_limit,