Skip to content

Commit 92766ad

Browse files
morishwemorishwe
committed
Adding chat history to Agent Chat. Adding a new table for session ID tracking
1 parent 2d42559 commit 92766ad

File tree

10 files changed

+1244
-82
lines changed

10 files changed

+1244
-82
lines changed

‎src/api/schema.graphql‎

Lines changed: 103 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
interface DynamoDbBase {
2-
PK: ID!
3-
SK: ID!
4-
ExpiresAfter: AWSTimestamp
2+
PK: ID!
3+
SK: ID!
4+
ExpiresAfter: AWSTimestamp
55
}
66

77
type Document implements DynamoDbBase @aws_cognito_user_pools @aws_iam {
@@ -30,7 +30,7 @@ type Document implements DynamoDbBase @aws_cognito_user_pools @aws_iam {
3030

3131
type Section @aws_cognito_user_pools @aws_iam {
3232
Id: String
33-
PageIds: [ Int ]
33+
PageIds: [Int]
3434
Class: String
3535
OutputJSONUri: String
3636
ConfidenceThresholdAlerts: [ConfidenceThresholdAlert]
@@ -51,18 +51,18 @@ type Page @aws_cognito_user_pools @aws_iam {
5151
}
5252

5353
type DocumentList @aws_cognito_user_pools @aws_iam {
54-
Documents: [DocumentListItem]
55-
nextToken: String
54+
Documents: [DocumentListItem]
55+
nextToken: String
5656
}
5757

5858
type DocumentListItem implements DynamoDbBase @aws_cognito_user_pools @aws_iam {
59-
PK: ID!
60-
SK: ID!
61-
ObjectKey: ID
62-
InitialEventTime: AWSDateTime
63-
ExpiresAfter: AWSTimestamp
64-
HITLStatus: String
65-
HITLReviewURL: String
59+
PK: ID!
60+
SK: ID!
61+
ObjectKey: ID
62+
InitialEventTime: AWSDateTime
63+
ExpiresAfter: AWSTimestamp
64+
HITLStatus: String
65+
HITLReviewURL: String
6666
}
6767

6868
type ValidationError @aws_cognito_user_pools @aws_iam {
@@ -104,7 +104,6 @@ type DiscoveryJobListItem @aws_cognito_user_pools @aws_iam {
104104
createdAt: String
105105
updatedAt: String
106106
errorMessage: String
107-
108107
}
109108

110109
input CreateDocumentInput {
@@ -155,7 +154,7 @@ input UpdateDocumentInput {
155154

156155
input SectionInput {
157156
Id: String
158-
PageIds: [ Int ]
157+
PageIds: [Int]
159158
Class: String
160159
OutputJSONUri: String
161160
ConfidenceThresholdAlerts: [ConfidenceThresholdAlertInput]
@@ -225,7 +224,7 @@ type StepFunctionExecutionResponse @aws_cognito_user_pools @aws_iam {
225224
steps: [StepFunctionExecutionStep]
226225
}
227226

228-
type DiscoveryJob @aws_cognito_user_pools @aws_iam {
227+
type DiscoveryJob @aws_cognito_user_pools @aws_iam {
229228
jobId: ID!
230229
status: String!
231230
errorMessage: String
@@ -240,14 +239,14 @@ type Agent @aws_cognito_user_pools @aws_iam {
240239

241240
type AgentJob @aws_cognito_user_pools @aws_iam {
242241
jobId: ID!
243-
status: String! # "PENDING", "PROCESSING", "COMPLETED", "FAILED"
242+
status: String! # "PENDING", "PROCESSING", "COMPLETED", "FAILED"
244243
query: String!
245-
agentIds: String # JSON string containing list of agent IDs used
244+
agentIds: String # JSON string containing list of agent IDs used
246245
createdAt: AWSDateTime!
247246
completedAt: AWSDateTime
248247
result: String
249248
error: String
250-
agent_messages: String # JSON string containing agent conversation messages
249+
agent_messages: String # JSON string containing agent conversation messages
251250
}
252251

253252
type AgentJobConnection @aws_cognito_user_pools @aws_iam {
@@ -263,6 +262,20 @@ type AgentChatMessage @aws_cognito_user_pools @aws_iam {
263262
sessionId: String
264263
}
265264

265+
type ChatSession @aws_cognito_user_pools @aws_iam {
266+
sessionId: ID!
267+
title: String!
268+
createdAt: AWSDateTime!
269+
updatedAt: AWSDateTime!
270+
messageCount: Int!
271+
lastMessage: String
272+
}
273+
274+
type ChatSessionConnection @aws_cognito_user_pools @aws_iam {
275+
items: [ChatSession]
276+
nextToken: String
277+
}
278+
266279
type MessageContent {
267280
text: String
268281
}
@@ -271,42 +284,95 @@ type Mutation {
271284
createDocument(input: CreateDocumentInput!): CreateDocumentOutput @aws_iam
272285
updateDocument(input: UpdateDocumentInput!): Document @aws_iam
273286
deleteDocument(objectKeys: [String!]!): Boolean! @aws_cognito_user_pools
274-
updateConfiguration(customConfig: AWSJSON!): UpdateConfigurationResponse @aws_cognito_user_pools
275-
uploadDocument(fileName: String!, contentType: String, prefix: String, bucket: String): PresignedUrlResponse! @aws_cognito_user_pools
276-
uploadDiscoveryDocument(fileName: String!, contentType: String, prefix: String, bucket: String, groundTruthFileName: String): DisPresignedUrlResponse! @aws_cognito_user_pools
277-
copyToBaseline(objectKey: String!): CopyToBaselineResponse! @aws_cognito_user_pools
287+
updateConfiguration(customConfig: AWSJSON!): UpdateConfigurationResponse
288+
@aws_cognito_user_pools
289+
uploadDocument(
290+
fileName: String!
291+
contentType: String
292+
prefix: String
293+
bucket: String
294+
): PresignedUrlResponse! @aws_cognito_user_pools
295+
uploadDiscoveryDocument(
296+
fileName: String!
297+
contentType: String
298+
prefix: String
299+
bucket: String
300+
groundTruthFileName: String
301+
): DisPresignedUrlResponse! @aws_cognito_user_pools
302+
copyToBaseline(objectKey: String!): CopyToBaselineResponse!
303+
@aws_cognito_user_pools
278304
reprocessDocument(objectKeys: [String!]!): Boolean! @aws_cognito_user_pools
279-
processChanges(objectKey: String!, modifiedSections: [ModifiedSectionInput!]!): ProcessChangesResponse! @aws_cognito_user_pools
280-
updateAgentJobStatus(jobId: ID!, status: String!, userId: String!, result: String): Boolean @aws_iam
281-
updateDiscoveryJobStatus(jobId: ID!, status: String!, errorMessage: String): DiscoveryJob @aws_iam
305+
processChanges(
306+
objectKey: String!
307+
modifiedSections: [ModifiedSectionInput!]!
308+
): ProcessChangesResponse! @aws_cognito_user_pools
309+
updateAgentJobStatus(
310+
jobId: ID!
311+
status: String!
312+
userId: String!
313+
result: String
314+
): Boolean @aws_iam
315+
updateDiscoveryJobStatus(
316+
jobId: ID!
317+
status: String!
318+
errorMessage: String
319+
): DiscoveryJob @aws_iam
282320
deleteAgentJob(jobId: ID!): Boolean @aws_cognito_user_pools
283-
sendAgentChatMessage(prompt: String!, sessionId: String, method: String, enableCodeIntelligence: Boolean): AgentChatMessage @aws_cognito_user_pools @aws_iam
284-
updateAgentChatMessage(sessionId: ID!, content: String!): AgentChatMessage @aws_cognito_user_pools @aws_iam
321+
sendAgentChatMessage(
322+
prompt: String!
323+
sessionId: String
324+
method: String
325+
enableCodeIntelligence: Boolean
326+
): AgentChatMessage @aws_cognito_user_pools @aws_iam
327+
updateAgentChatMessage(sessionId: ID!, content: String!): AgentChatMessage
328+
@aws_cognito_user_pools
329+
@aws_iam
330+
deleteChatSession(sessionId: ID!): Boolean @aws_cognito_user_pools @aws_iam
331+
updateChatSessionTitle(sessionId: ID!, title: String!): ChatSession
332+
@aws_cognito_user_pools
333+
@aws_iam
285334
}
286335

287336
type Query @aws_cognito_user_pools @aws_iam {
288337
getDocument(ObjectKey: ID!): Document
289-
listDocuments(startDateTime: AWSDateTime, endDateTime: AWSDateTime): DocumentList
338+
listDocuments(
339+
startDateTime: AWSDateTime
340+
endDateTime: AWSDateTime
341+
): DocumentList
290342
listDocumentsDateHour(date: AWSDate, hour: Int): DocumentList
291-
listDocumentsDateShard(date: AWSDate, shard: Int): DocumentList
343+
listDocumentsDateShard(date: AWSDate, shard: Int): DocumentList
292344
getFileContents(s3Uri: String!): FileContentsResponse
293345
getConfiguration: ConfigurationResponse
294346
listDiscoveryJobs: DiscoveryJobList
295347
queryKnowledgeBase(input: String!, sessionId: String): String
296-
chatWithDocument(s3Uri: String!, prompt: String!, history: AWSJSON!, modelId: String!): String
348+
chatWithDocument(
349+
s3Uri: String!
350+
prompt: String!
351+
history: AWSJSON!
352+
modelId: String!
353+
): String
297354
getStepFunctionExecution(executionArn: String!): StepFunctionExecutionResponse
298355
listAvailableAgents: [Agent] @aws_cognito_user_pools @aws_iam
299-
submitAgentQuery(query: String!, agentIds: [String!]!): AgentJob @aws_cognito_user_pools
356+
submitAgentQuery(query: String!, agentIds: [String!]!): AgentJob
357+
@aws_cognito_user_pools
300358
getAgentJobStatus(jobId: ID!): AgentJob @aws_cognito_user_pools
301-
listAgentJobs(limit: Int, nextToken: String): AgentJobConnection @aws_cognito_user_pools
302-
getAgentChatMessages(sessionId: ID!): [AgentChatMessage] @aws_cognito_user_pools @aws_iam
359+
listAgentJobs(limit: Int, nextToken: String): AgentJobConnection
360+
@aws_cognito_user_pools
361+
getAgentChatMessages(sessionId: ID!): [AgentChatMessage]
362+
@aws_cognito_user_pools
363+
@aws_iam
364+
listChatSessions(limit: Int, nextToken: String): ChatSessionConnection
365+
@aws_cognito_user_pools
366+
@aws_iam
367+
getChatMessages(sessionId: ID!): [AgentChatMessage]
368+
@aws_cognito_user_pools
369+
@aws_iam
303370
}
304371

305372
type Subscription @aws_cognito_user_pools @aws_iam {
306373
onCreateDocument: CreateDocumentOutput
307-
@aws_subscribe(mutations: ["createDocument"])
308-
onUpdateDocument: Document
309-
@aws_subscribe(mutations: ["updateDocument"])
374+
@aws_subscribe(mutations: ["createDocument"])
375+
onUpdateDocument: Document @aws_subscribe(mutations: ["updateDocument"])
310376
onAgentJobComplete(jobId: ID!): Boolean
311377
@aws_subscribe(mutations: ["updateAgentJobStatus"])
312378
onDiscoveryJobStatusChange(jobId: ID!): DiscoveryJob

‎src/lambda/agent_chat_resolver/index.py‎

Lines changed: 105 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626

2727
# Get environment variables
2828
CHAT_MESSAGES_TABLE = os.environ.get("CHAT_MESSAGES_TABLE")
29+
CHAT_SESSIONS_TABLE = os.environ.get("CHAT_SESSIONS_TABLE")
2930
AGENT_CHAT_PROCESSOR_FUNCTION = os.environ.get("AGENT_CHAT_PROCESSOR_FUNCTION")
3031
DATA_RETENTION_DAYS = int(os.environ.get("DATA_RETENTION_DAYS", "30"))
3132

@@ -120,6 +121,14 @@ def handler(event, context):
120121

121122
table.put_item(Item=message)
122123
logger.info(f"Stored message in DynamoDB for session {session_id}: {method}")
124+
125+
# Manage session metadata for user messages
126+
if not is_assistant_response: # This is a user message
127+
# Get user identity for session metadata
128+
identity = event.get("identity", {})
129+
user_id = identity.get("username") or identity.get("sub") or "anonymous"
130+
131+
manage_session_metadata(user_id, session_id, prompt, timestamp, expires_after)
123132
else:
124133
logger.info(f"Skipped storing streaming message in DynamoDB: {method}")
125134

@@ -159,24 +168,101 @@ def handler(event, context):
159168
"isProcessing": False,
160169
"sessionId": str(session_id) if session_id else ""
161170
}
171+
172+
173+
def manage_session_metadata(user_id, session_id, prompt, timestamp, expires_after):
174+
"""
175+
Create or update session metadata in the ChatSessionsTable.
176+
177+
Args:
178+
user_id: The user ID for the session
179+
session_id: The session ID
180+
prompt: The user's message content
181+
timestamp: The message timestamp
182+
expires_after: TTL for the session
183+
"""
184+
try:
185+
if not CHAT_SESSIONS_TABLE:
186+
logger.warn("CHAT_SESSIONS_TABLE not configured, skipping session metadata management")
187+
return
188+
189+
sessions_table = dynamodb.Table(CHAT_SESSIONS_TABLE)
190+
191+
# Check if session already exists
192+
try:
193+
response = sessions_table.get_item(
194+
Key={
195+
"userId": user_id,
196+
"sessionId": session_id
197+
}
198+
)
199+
200+
if response.get("Item"):
201+
# Session exists, update it
202+
update_session_metadata(sessions_table, user_id, session_id, prompt, timestamp)
203+
else:
204+
# New session, create it
205+
create_session_metadata(sessions_table, user_id, session_id, prompt, timestamp, expires_after)
206+
207+
except ClientError as e:
208+
logger.error(f"Error managing session metadata: {str(e)}")
209+
# Don't fail the main operation if session metadata fails
210+
162211
except Exception as e:
163-
# Check if this is a validation error that should propagate as GraphQL error
164-
error_str = str(e)
165-
if any(msg in error_str for msg in [
166-
"prompt parameter is required",
167-
"sessionId parameter is required",
168-
"Prompt exceeds maximum length"
169-
]):
170-
# Re-raise validation errors so they become GraphQL errors
171-
raise e
172-
173-
# Handle other unexpected errors
174-
error_msg = f"Error processing request: {error_str}"
175-
logger.error(error_msg)
176-
return {
177-
"role": "assistant",
178-
"content": f"Error: {error_msg}",
179-
"timestamp": datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
180-
"isProcessing": False,
181-
"sessionId": str(session_id) if session_id else ""
212+
logger.error(f"Unexpected error managing session metadata: {str(e)}")
213+
# Don't fail the main operation if session metadata fails
214+
215+
216+
def create_session_metadata(sessions_table, user_id, session_id, prompt, timestamp, expires_after):
217+
"""Create a new session metadata record."""
218+
try:
219+
# Generate session title from the first message
220+
title = prompt.strip()
221+
if len(title) > 50:
222+
title = title[:47] + "..."
223+
224+
# Create last message preview
225+
last_message = prompt[:100] + "..." if len(prompt) > 100 else prompt
226+
227+
session_record = {
228+
"userId": user_id,
229+
"sessionId": session_id,
230+
"title": title,
231+
"createdAt": timestamp,
232+
"updatedAt": timestamp,
233+
"messageCount": 1, # First user message
234+
"lastMessage": last_message,
235+
"ExpiresAfter": expires_after
182236
}
237+
238+
sessions_table.put_item(Item=session_record)
239+
logger.info(f"Created session metadata for {session_id}")
240+
241+
except ClientError as e:
242+
logger.error(f"Error creating session metadata: {str(e)}")
243+
244+
245+
def update_session_metadata(sessions_table, user_id, session_id, prompt, timestamp):
246+
"""Update existing session metadata record."""
247+
try:
248+
# Create last message preview
249+
last_message = prompt[:100] + "..." if len(prompt) > 100 else prompt
250+
251+
# Update the session record
252+
sessions_table.update_item(
253+
Key={
254+
"userId": user_id,
255+
"sessionId": session_id
256+
},
257+
UpdateExpression="SET updatedAt = :timestamp, messageCount = messageCount + :inc, lastMessage = :last_message",
258+
ExpressionAttributeValues={
259+
":timestamp": timestamp,
260+
":inc": 1,
261+
":last_message": last_message
262+
}
263+
)
264+
265+
logger.info(f"Updated session metadata for {session_id}")
266+
267+
except ClientError as e:
268+
logger.error(f"Error updating session metadata: {str(e)}")

0 commit comments

Comments
 (0)