Skip to content

Commit 2b2557a

Browse files
MikeAlhayekMikeAlhayek
authored
Ensure the orchestrator can access and track chat-session-level documents (#353)
1 parent 18ae3ba commit 2b2557a

File tree

26 files changed

+295
-120
lines changed

26 files changed

+295
-120
lines changed

‎src/Abstractions/CrestApps.OrchardCore.AI.Abstractions/Models/AIChatSession.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ public sealed class AIChatSession : Entity
3838
/// Gets or sets the collection of document references attached to this session.
3939
/// Documents are uploaded by users and used for RAG (Retrieval-Augmented Generation).
4040
/// </summary>
41-
public IList<ChatInteractionDocumentInfo> Documents { get; set; } = [];
41+
public List<ChatDocumentInfo> Documents { get; set; } = [];
4242

4343
/// <summary>
4444
/// Gets or sets the UTC date and time when the session was first created.

src/Abstractions/CrestApps.OrchardCore.AI.Abstractions/Models/ChatInteractionDocumentInfo.cs renamed to src/Abstractions/CrestApps.OrchardCore.AI.Abstractions/Models/ChatDocumentInfo.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
namespace CrestApps.OrchardCore.AI.Models;
22

3-
public sealed class ChatInteractionDocumentInfo
3+
public sealed class ChatDocumentInfo
44
{
55
/// <summary>
66
/// Gets or sets the unique identifier for this document.

‎src/Abstractions/CrestApps.OrchardCore.AI.Abstractions/Models/ChatInteraction.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ public sealed class ChatInteraction : CatalogItem, ISourceAwareModel
9999
/// Gets or sets the collection of attached documents for "chat against own data" functionality.
100100
/// Only applicable when Source is AzureOpenAIOwnData.
101101
/// </summary>
102-
public IList<ChatInteractionDocumentInfo> Documents { get; set; } = [];
102+
public List<ChatDocumentInfo> Documents { get; set; } = [];
103103

104104
/// <summary>
105105
/// Gets or sets the UTC date and time when the interaction was created.

‎src/Abstractions/CrestApps.OrchardCore.AI.Abstractions/Models/OrchestratorContext.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ public sealed class OrchestrationContext
4848
/// Gets or sets the document references available for this session.
4949
/// System tools use these to load document content on demand.
5050
/// </summary>
51-
public IList<ChatInteractionDocumentInfo> Documents { get; set; } = [];
51+
public List<ChatDocumentInfo> Documents { get; set; } = [];
5252

5353
/// <summary>
5454
/// Gets or sets whether tools should be disabled for this orchestration.

‎src/Core/CrestApps.OrchardCore.AI.Chat.Interactions.Core/Tools/ListDocumentsTool.cs‎

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -78,14 +78,34 @@ protected override async ValueTask<object> InvokeCoreAsync(
7878
return "Document store is not available.";
7979
}
8080

81-
var documents = await documentStore.GetDocumentsAsync(profile.ItemId, AIConstants.DocumentReferenceTypes.Profile);
81+
// Collect documents from both profile-level and session-level sources.
82+
var allDocuments = new List<AIDocument>();
8283

83-
if (documents is null || documents.Count == 0)
84+
var profileDocs = await documentStore.GetDocumentsAsync(profile.ItemId, AIConstants.DocumentReferenceTypes.Profile);
85+
86+
if (profileDocs is { Count: > 0 })
8487
{
85-
return "No documents are attached to this profile.";
88+
allDocuments.AddRange(profileDocs);
8689
}
8790

88-
var result = documents.Select(d => new
91+
if (AIInvocationScope.Current?.Items.TryGetValue(nameof(AIChatSession), out var sessionObj) == true &&
92+
sessionObj is AIChatSession session &&
93+
session.Documents is { Count: > 0 })
94+
{
95+
var sessionDocs = await documentStore.GetDocumentsAsync(session.SessionId, AIConstants.DocumentReferenceTypes.ChatSession);
96+
97+
if (sessionDocs is { Count: > 0 })
98+
{
99+
allDocuments.AddRange(sessionDocs);
100+
}
101+
}
102+
103+
if (allDocuments.Count == 0)
104+
{
105+
return "No documents are attached.";
106+
}
107+
108+
var result = allDocuments.Select(d => new
89109
{
90110
d.ItemId,
91111
d.FileName,

‎src/Core/CrestApps.OrchardCore.AI.Chat.Interactions.Core/Tools/ReadDocumentTool.cs‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,11 +80,24 @@ protected override async ValueTask<object> InvokeCoreAsync(
8080
return "Document store is not available.";
8181
}
8282

83+
// The document could belong to either the profile or a chat session.
84+
var validReferenceIds = new HashSet<string>(StringComparer.OrdinalIgnoreCase)
85+
{
86+
profile.ItemId,
87+
};
88+
89+
if (AIInvocationScope.Current?.Items.TryGetValue(nameof(AIChatSession), out var sessionObj) == true &&
90+
sessionObj is AIChatSession session &&
91+
session.Documents is { Count: > 0 })
92+
{
93+
validReferenceIds.Add(session.SessionId);
94+
}
95+
8396
var document = await documentStore.FindByIdAsync(documentId);
8497

85-
if (document is null || document.ReferenceId != profile.ItemId)
98+
if (document is null || !validReferenceIds.Contains(document.ReferenceId))
8699
{
87-
return $"Document with ID '{documentId}' was not found in this profile.";
100+
return $"Document with ID '{documentId}' was not found.";
88101
}
89102

90103
return FormatDocumentText(document.FileName, document.Text);

‎src/Core/CrestApps.OrchardCore.AI.Chat.Interactions.Core/Tools/ReadTabularDataTool.cs‎

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -69,25 +69,49 @@ protected override async ValueTask<object> InvokeCoreAsync(
6969

7070
var executionContext = AIInvocationScope.Current?.ToolExecutionContext;
7171

72-
if (executionContext?.Resource is not ChatInteraction interaction)
72+
string referenceId = null;
73+
HashSet<string> validReferenceIds = null;
74+
75+
if (executionContext?.Resource is ChatInteraction interaction)
76+
{
77+
referenceId = interaction.ItemId;
78+
}
79+
else if (executionContext?.Resource is AIProfile profile)
80+
{
81+
referenceId = profile.ItemId;
82+
validReferenceIds = new HashSet<string>(StringComparer.OrdinalIgnoreCase)
83+
{
84+
profile.ItemId,
85+
};
86+
87+
// Documents could be attached at the profile or the session level.
88+
if (AIInvocationScope.Current?.Items.TryGetValue(nameof(AIChatSession), out var sessionObj) == true &&
89+
sessionObj is AIChatSession session &&
90+
session.Documents is { Count: > 0 })
91+
{
92+
validReferenceIds.Add(session.SessionId);
93+
}
94+
}
95+
96+
if (string.IsNullOrEmpty(referenceId))
7397
{
74-
return "Document access requires an active chat interaction session.";
98+
return "Document access requires an active chat interaction session or AI profile.";
7599
}
76100

77-
var chatInteractionId = interaction.ItemId;
78101
var documentStore = arguments.Services.GetService<IAIDocumentStore>();
79102

80103
if (documentStore is null)
81104
{
82105
return "Document store is not available.";
83106
}
84107

85-
// Query only documents belonging to this interaction to prevent cross-session access.
108+
// Query only documents belonging to this resource to prevent cross-session access.
86109
var document = await documentStore.FindByIdAsync(documentId);
87110

88-
if (document is null || document.ReferenceId != chatInteractionId)
111+
if (document is null ||
112+
(validReferenceIds is not null ? !validReferenceIds.Contains(document.ReferenceId) : document.ReferenceId != referenceId))
89113
{
90-
return $"Document with ID '{documentId}' was not found in this session.";
114+
return $"Document with ID '{documentId}' was not found.";
91115
}
92116

93117
if (string.IsNullOrWhiteSpace(document.Text))

‎src/Core/CrestApps.OrchardCore.AI.Core/Handlers/DocumentOrchestrationHandler.cs‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,15 +35,17 @@ public Task BuildingAsync(OrchestrationContextBuildingContext context)
3535
if (context.Resource is ChatInteraction interaction &&
3636
interaction.Documents is { Count: > 0 })
3737
{
38-
context.Context.Documents = interaction.Documents;
38+
context.Context.Documents ??= [];
39+
context.Context.Documents.AddRange(interaction.Documents);
3940
}
4041
else if (context.Resource is AIProfile profile)
4142
{
4243
var documentsMetadata = profile.As<DocumentsMetadata>();
4344

4445
if (documentsMetadata.Documents is { Count: > 0 })
4546
{
46-
context.Context.Documents = documentsMetadata.Documents;
47+
context.Context.Documents ??= [];
48+
context.Context.Documents.AddRange(documentsMetadata.Documents);
4749
}
4850
}
4951

@@ -62,7 +64,8 @@ context.Resource is AIProfile &&
6264
sessionObj is AIChatSession session &&
6365
session.Documents is { Count: > 0 })
6466
{
65-
context.OrchestrationContext.Documents = session.Documents;
67+
context.OrchestrationContext.Documents ??= [];
68+
context.OrchestrationContext.Documents.AddRange(session.Documents);
6669
}
6770

6871
if (context.OrchestrationContext.Documents is not { Count: > 0 } ||

‎src/Core/CrestApps.OrchardCore.AI.Core/Handlers/PreemptiveRagOrchestrationHandler.cs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,7 @@ public async Task BuiltAsync(OrchestrationContextBuiltContext buildContext)
115115

116116
if (!string.IsNullOrEmpty(prompt))
117117
{
118+
buildContext.OrchestrationContext.SystemMessageBuilder.AppendLine();
118119
buildContext.OrchestrationContext.SystemMessageBuilder.AppendLine(prompt);
119120
}
120121
}
@@ -133,6 +134,7 @@ public async Task BuiltAsync(OrchestrationContextBuiltContext buildContext)
133134

134135
if (!string.IsNullOrEmpty(prompt))
135136
{
137+
buildContext.OrchestrationContext.SystemMessageBuilder.AppendLine();
136138
buildContext.OrchestrationContext.SystemMessageBuilder.AppendLine(prompt);
137139
}
138140
}
@@ -142,6 +144,7 @@ public async Task BuiltAsync(OrchestrationContextBuiltContext buildContext)
142144

143145
if (!string.IsNullOrEmpty(prompt))
144146
{
147+
buildContext.OrchestrationContext.SystemMessageBuilder.AppendLine();
145148
buildContext.OrchestrationContext.SystemMessageBuilder.AppendLine(prompt);
146149
}
147150
}
@@ -152,6 +155,7 @@ public async Task BuiltAsync(OrchestrationContextBuiltContext buildContext)
152155

153156
if (!string.IsNullOrEmpty(prompt))
154157
{
158+
buildContext.OrchestrationContext.SystemMessageBuilder.AppendLine();
155159
buildContext.OrchestrationContext.SystemMessageBuilder.AppendLine(prompt);
156160
}
157161
}
@@ -170,6 +174,7 @@ private async Task InjectToolSearchInstructionsAsync(OrchestrationContextBuiltCo
170174

171175
if (!string.IsNullOrEmpty(prompt))
172176
{
177+
buildContext.OrchestrationContext.SystemMessageBuilder.AppendLine();
173178
buildContext.OrchestrationContext.SystemMessageBuilder.AppendLine(prompt);
174179
}
175180
}
@@ -180,6 +185,7 @@ private async Task InjectToolSearchInstructionsAsync(OrchestrationContextBuiltCo
180185

181186
if (!string.IsNullOrEmpty(prompt))
182187
{
188+
buildContext.OrchestrationContext.SystemMessageBuilder.AppendLine();
183189
buildContext.OrchestrationContext.SystemMessageBuilder.AppendLine(prompt);
184190
}
185191
}

‎src/Core/CrestApps.OrchardCore.AI.Core/Models/DocumentsMetadata.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ public sealed class DocumentsMetadata
99
/// <summary>
1010
/// Gets or sets the collection of attached document metadata.
1111
/// </summary>
12-
public IList<ChatInteractionDocumentInfo> Documents { get; set; } = [];
12+
public IList<ChatDocumentInfo> Documents { get; set; } = [];
1313

1414
/// <summary>
1515
/// Gets or sets the number of top matching document chunks to include in AI context.

0 commit comments

Comments
 (0)