From 6307c7bb28ca4c1c568f5b2e2f4a78e9b6c420c5 Mon Sep 17 00:00:00 2001 From: James Brunton Date: Tue, 21 Apr 2026 16:03:10 +0100 Subject: [PATCH] Allow chat history to be sent to AI engine (#6128) # Description of Changes Add an extra parameter to every agent to receive the conversation history in addition to the current message. This will make it possible to answer followup questions from the AI without needing to give full context in your message. --- .../model/api/ai/AiConversationMessage.java | 22 ++++++++++++++++++ .../model/api/ai/AiWorkflowRequest.java | 7 ++++++ .../service/AiWorkflowService.java | 7 ++++++ engine/src/stirling/agents/orchestrator.py | 23 ++++++++++++++++--- engine/src/stirling/agents/pdf_edit.py | 2 ++ engine/src/stirling/agents/pdf_questions.py | 9 +++++++- engine/src/stirling/agents/user_spec.py | 20 ++++++++-------- engine/src/stirling/contracts/__init__.py | 2 ++ engine/src/stirling/contracts/common.py | 6 +++++ engine/src/stirling/contracts/orchestrator.py | 9 +++++++- engine/src/stirling/contracts/pdf_edit.py | 3 ++- .../src/stirling/contracts/pdf_questions.py | 9 +++++++- engine/tests/test_user_spec_agent.py | 8 +++++-- .../components/chat/ChatContext.tsx | 8 +++++++ 14 files changed, 116 insertions(+), 19 deletions(-) create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiConversationMessage.java diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiConversationMessage.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiConversationMessage.java new file mode 100644 index 0000000000..32a811302c --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiConversationMessage.java @@ -0,0 +1,22 @@ +package stirling.software.proprietary.model.api.ai; + +import io.swagger.v3.oas.annotations.media.Schema; + +import jakarta.validation.constraints.NotBlank; +import jakarta.validation.constraints.NotNull; + +import lombok.Data; + +@Data +@Schema(description = "A prior message in the chat conversation") +public class AiConversationMessage { + + @NotNull + @NotBlank + @Schema(description = "The role of the message sender", example = "user") + private String role; + + @NotNull + @Schema(description = "The content of the message") + private String content; +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowRequest.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowRequest.java index 22228d2aa3..da327177c4 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowRequest.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowRequest.java @@ -1,5 +1,6 @@ package stirling.software.proprietary.model.api.ai; +import java.util.ArrayList; import java.util.List; import io.swagger.v3.oas.annotations.media.Schema; @@ -20,4 +21,10 @@ public class AiWorkflowRequest { @NotBlank @Schema(description = "The user message to orchestrate", example = "Summarise these documents") private String userMessage; + + @Schema( + description = + "Prior chat messages exchanged between the user and the assistant, ordered" + + " oldest-first. Excludes the current userMessage.") + private List conversationHistory = new ArrayList<>(); } diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java index 0d02a44381..14bc56ed98 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java @@ -34,6 +34,7 @@ import stirling.software.common.util.ExceptionUtils; import stirling.software.common.util.TempFile; import stirling.software.common.util.TempFileManager; import stirling.software.common.util.ZipExtractionUtils; +import stirling.software.proprietary.model.api.ai.AiConversationMessage; import stirling.software.proprietary.model.api.ai.AiWorkflowFileInput; import stirling.software.proprietary.model.api.ai.AiWorkflowFileRequest; import stirling.software.proprietary.model.api.ai.AiWorkflowOutcome; @@ -92,6 +93,10 @@ public class AiWorkflowService { WorkflowTurnRequest initialRequest = new WorkflowTurnRequest(); initialRequest.setUserMessage(request.getUserMessage().trim()); initialRequest.setFileNames(new ArrayList<>(filesByName.keySet())); + initialRequest.setConversationHistory( + request.getConversationHistory() == null + ? new ArrayList<>() + : new ArrayList<>(request.getConversationHistory())); listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.ANALYZING)); @@ -183,6 +188,7 @@ public class AiWorkflowService { WorkflowTurnRequest nextRequest = new WorkflowTurnRequest(); nextRequest.setUserMessage(request.getUserMessage()); nextRequest.setFileNames(request.getFileNames()); + nextRequest.setConversationHistory(request.getConversationHistory()); nextRequest.setArtifacts(pdfContentExtractor.buildArtifacts(contentResults)); nextRequest.setResumeWith(response.getResumeWith()); return new WorkflowState.Pending(nextRequest); @@ -415,6 +421,7 @@ public class AiWorkflowService { private static class WorkflowTurnRequest { private String userMessage; private List fileNames = new ArrayList<>(); + private List conversationHistory = new ArrayList<>(); private List artifacts = new ArrayList<>(); private String resumeWith; } diff --git a/engine/src/stirling/agents/orchestrator.py b/engine/src/stirling/agents/orchestrator.py index e4c17d0059..c2537f78ce 100644 --- a/engine/src/stirling/agents/orchestrator.py +++ b/engine/src/stirling/agents/orchestrator.py @@ -23,6 +23,7 @@ from stirling.contracts import ( SupportedCapability, ToolOperationStep, UnsupportedCapabilityResponse, + format_conversation_history, ) from stirling.contracts.pdf_edit import EditPlanResponse from stirling.models.agent_tool_models import AgentToolId, MathAuditorAgentParams @@ -117,7 +118,11 @@ class OrchestratorAgent: async def _run_pdf_edit(self, request: OrchestratorRequest) -> PdfEditResponse: return await PdfEditAgent(self.runtime).handle( - PdfEditRequest(user_message=request.user_message, file_names=request.file_names) + PdfEditRequest( + user_message=request.user_message, + file_names=request.file_names, + conversation_history=request.conversation_history, + ) ) async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionResponse: @@ -130,6 +135,7 @@ class OrchestratorAgent: question=request.user_message, file_names=request.file_names, page_text=extracted_text.files if extracted_text is not None else [], + conversation_history=request.conversation_history, ) ) @@ -137,7 +143,12 @@ class OrchestratorAgent: return await self._run_agent_draft(ctx.deps.request) async def _run_agent_draft(self, request: OrchestratorRequest) -> AgentDraftWorkflowResponse: - return await UserSpecAgent(self.runtime).draft(AgentDraftRequest(user_message=request.user_message)) + return await UserSpecAgent(self.runtime).draft( + AgentDraftRequest( + user_message=request.user_message, + conversation_history=request.conversation_history, + ) + ) async def math_auditor_agent(self, ctx: RunContext[OrchestratorDeps]) -> EditPlanResponse: return EditPlanResponse( @@ -167,7 +178,13 @@ class OrchestratorAgent: def _build_prompt(self, request: OrchestratorRequest) -> str: artifact_summary = self._describe_artifacts(request) file_names = ", ".join(request.file_names) if request.file_names else "Unknown files" - return f"User message: {request.user_message}\nFiles: {file_names}\nAvailable artifacts:\n{artifact_summary}" + history = format_conversation_history(request.conversation_history) + return ( + f"Conversation history:\n{history}\n" + f"User message: {request.user_message}\n" + f"Files: {file_names}\n" + f"Available artifacts:\n{artifact_summary}" + ) def _describe_artifacts(self, request: OrchestratorRequest) -> str: if not request.artifacts: diff --git a/engine/src/stirling/agents/pdf_edit.py b/engine/src/stirling/agents/pdf_edit.py index b5ad8c431d..40b1365783 100644 --- a/engine/src/stirling/agents/pdf_edit.py +++ b/engine/src/stirling/agents/pdf_edit.py @@ -13,6 +13,7 @@ from stirling.contracts import ( PdfEditRequest, PdfEditResponse, ToolOperationStep, + format_conversation_history, ) from stirling.models import OPERATIONS, ApiModel, ParamToolModel, ToolEndpoint from stirling.services import AppRuntime @@ -146,6 +147,7 @@ class PdfEditAgent: def _build_selection_prompt(self, request: PdfEditRequest) -> str: file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided." return ( + f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n" f"User request: {request.user_message}\n" f"Files: {file_names}\n" f"Supported operations: {self._supported_operations_prompt()}\n" diff --git a/engine/src/stirling/agents/pdf_questions.py b/engine/src/stirling/agents/pdf_questions.py index 60f7ea94c2..c646159da8 100644 --- a/engine/src/stirling/agents/pdf_questions.py +++ b/engine/src/stirling/agents/pdf_questions.py @@ -12,6 +12,7 @@ from stirling.contracts import ( PdfQuestionNotFoundResponse, PdfQuestionRequest, PdfQuestionResponse, + format_conversation_history, ) from stirling.services import AppRuntime @@ -70,7 +71,13 @@ class PdfQuestionAgent: for selection in file_text.pages ] pages = "\n\n".join(sections) - return f"Files: {file_names}\nQuestion: {request.question}\nExtracted page text:\n{pages}" + history = format_conversation_history(request.conversation_history) + return ( + f"Conversation history:\n{history}\n" + f"Files: {file_names}\n" + f"Question: {request.question}\n" + f"Extracted page text:\n{pages}" + ) def _has_page_text(self, page_text: list[ExtractedFileText]) -> bool: return any(selection.text.strip() for file_text in page_text for selection in file_text.pages) diff --git a/engine/src/stirling/agents/user_spec.py b/engine/src/stirling/agents/user_spec.py index d30b5e53d5..40a028aa1b 100644 --- a/engine/src/stirling/agents/user_spec.py +++ b/engine/src/stirling/agents/user_spec.py @@ -18,6 +18,7 @@ from stirling.contracts import ( EditClarificationRequest, EditPlanResponse, PdfEditRequest, + format_conversation_history, ) from stirling.models import ApiModel from stirling.services import AppRuntime @@ -45,14 +46,15 @@ class UserSpecAgent: ) async def draft(self, request: AgentDraftRequest) -> AgentDraftWorkflowResponse: - edit_plan = await self._build_edit_plan(request.user_message) + edit_plan = await self._build_edit_plan(request.user_message, request.conversation_history) if not isinstance(edit_plan, EditPlanResponse): return edit_plan return AgentDraftResponse(draft=await self._run_draft_agent(request, edit_plan)) async def revise(self, request: AgentRevisionRequest) -> AgentRevisionWorkflowResponse: edit_plan = await self._build_edit_plan( - f"Current objective: {request.current_draft.objective}\nRevision request: {request.user_message}" + f"Current objective: {request.current_draft.objective}\nRevision request: {request.user_message}", + request.conversation_history, ) if not isinstance(edit_plan, EditPlanResponse): return edit_plan @@ -80,7 +82,7 @@ class UserSpecAgent: def _build_draft_prompt(self, request: AgentDraftRequest, edit_plan: EditPlanResponse) -> str: return ( f"User request:\n{request.user_message}\n\n" - f"Conversation history:\n{self._format_conversation_history(request.conversation_history)}\n\n" + f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n\n" f"Edit plan summary:\n{edit_plan.summary}\n\n" f"Edit plan rationale:\n{edit_plan.rationale or 'None'}\n\n" f"Edit plan steps:\n{edit_plan.model_dump_json(indent=2)}" @@ -89,20 +91,18 @@ class UserSpecAgent: def _build_revision_prompt(self, request: AgentRevisionRequest, edit_plan: EditPlanResponse) -> str: return ( f"Revision request:\n{request.user_message}\n\n" - f"Conversation history:\n{self._format_conversation_history(request.conversation_history)}\n\n" + f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n\n" f"Current draft:\n{request.current_draft.model_dump_json(indent=2)}\n\n" f"Edit plan summary:\n{edit_plan.summary}\n\n" f"Edit plan rationale:\n{edit_plan.rationale or 'None'}\n\n" f"Edit plan steps:\n{edit_plan.model_dump_json(indent=2)}" ) - def _format_conversation_history(self, conversation_history: list[ConversationMessage]) -> str: - if not conversation_history: - return "None" - return "\n".join(f"- {message.role}: {message.content}" for message in conversation_history) - async def _build_edit_plan( self, user_message: str, + conversation_history: list[ConversationMessage], ) -> EditPlanResponse | EditClarificationRequest | EditCannotDoResponse: - return await self.pdf_edit_agent.handle(PdfEditRequest(user_message=user_message)) + return await self.pdf_edit_agent.handle( + PdfEditRequest(user_message=user_message, conversation_history=conversation_history) + ) diff --git a/engine/src/stirling/contracts/__init__.py b/engine/src/stirling/contracts/__init__.py index 21633c075e..6fc075b444 100644 --- a/engine/src/stirling/contracts/__init__.py +++ b/engine/src/stirling/contracts/__init__.py @@ -18,6 +18,7 @@ from .common import ( SupportedCapability, ToolOperationStep, WorkflowOutcome, + format_conversation_history, ) from .execution import ( AgentExecutionRequest, @@ -104,6 +105,7 @@ __all__ = [ "Folio", "FolioManifest", "FolioType", + "format_conversation_history", "HealthResponse", "NeedContentFileRequest", "NextExecutionAction", diff --git a/engine/src/stirling/contracts/common.py b/engine/src/stirling/contracts/common.py index 9ee8856ddb..50038d1503 100644 --- a/engine/src/stirling/contracts/common.py +++ b/engine/src/stirling/contracts/common.py @@ -91,6 +91,12 @@ class ConversationMessage(ApiModel): content: str +def format_conversation_history(conversation_history: list[ConversationMessage]) -> str: + if not conversation_history: + return "None" + return "\n".join(f"- {message.role}: {message.content}" for message in conversation_history) + + class PdfTextSelection(ApiModel): page_number: int | None = None text: str diff --git a/engine/src/stirling/contracts/orchestrator.py b/engine/src/stirling/contracts/orchestrator.py index 3fa788fe66..7598372888 100644 --- a/engine/src/stirling/contracts/orchestrator.py +++ b/engine/src/stirling/contracts/orchestrator.py @@ -7,7 +7,13 @@ from pydantic import Field from stirling.models import ApiModel from .agent_drafts import AgentDraftResponse -from .common import ArtifactKind, ExtractedFileText, SupportedCapability, WorkflowOutcome +from .common import ( + ArtifactKind, + ConversationMessage, + ExtractedFileText, + SupportedCapability, + WorkflowOutcome, +) from .execution import NextExecutionAction from .pdf_edit import PdfEditResponse from .pdf_questions import PdfQuestionResponse @@ -24,6 +30,7 @@ WorkflowArtifact = Annotated[ExtractedTextArtifact, Field(discriminator="kind")] class OrchestratorRequest(ApiModel): user_message: str file_names: list[str] + conversation_history: list[ConversationMessage] = Field(default_factory=list) artifacts: list[WorkflowArtifact] = Field(default_factory=list) resume_with: SupportedCapability | None = None diff --git a/engine/src/stirling/contracts/pdf_edit.py b/engine/src/stirling/contracts/pdf_edit.py index 2bcfe7ac6f..a7bcedfa9d 100644 --- a/engine/src/stirling/contracts/pdf_edit.py +++ b/engine/src/stirling/contracts/pdf_edit.py @@ -6,12 +6,13 @@ from pydantic import Field from stirling.models import ApiModel -from .common import ToolOperationStep, WorkflowOutcome +from .common import ConversationMessage, ToolOperationStep, WorkflowOutcome class PdfEditRequest(ApiModel): user_message: str file_names: list[str] = Field(default_factory=list) + conversation_history: list[ConversationMessage] = Field(default_factory=list) class EditPlanResponse(ApiModel): diff --git a/engine/src/stirling/contracts/pdf_questions.py b/engine/src/stirling/contracts/pdf_questions.py index 987dc9b595..4ee0d25966 100644 --- a/engine/src/stirling/contracts/pdf_questions.py +++ b/engine/src/stirling/contracts/pdf_questions.py @@ -6,13 +6,20 @@ from pydantic import Field from stirling.models import ApiModel -from .common import ExtractedFileText, PdfContentType, SupportedCapability, WorkflowOutcome +from .common import ( + ConversationMessage, + ExtractedFileText, + PdfContentType, + SupportedCapability, + WorkflowOutcome, +) class PdfQuestionRequest(ApiModel): question: str page_text: list[ExtractedFileText] = Field(default_factory=list) file_names: list[str] + conversation_history: list[ConversationMessage] = Field(default_factory=list) class PdfQuestionAnswerResponse(ApiModel): diff --git a/engine/tests/test_user_spec_agent.py b/engine/tests/test_user_spec_agent.py index 1f9de664c3..3a91fcf1fb 100644 --- a/engine/tests/test_user_spec_agent.py +++ b/engine/tests/test_user_spec_agent.py @@ -32,7 +32,9 @@ class StubUserSpecAgent(UserSpecAgent): ], ) - async def _build_edit_plan(self, user_message: str) -> EditPlanResponse: + async def _build_edit_plan( + self, user_message: str, conversation_history: list[ConversationMessage] + ) -> EditPlanResponse: return self.edit_plan async def _run_draft_agent(self, request: AgentDraftRequest, edit_plan: EditPlanResponse) -> AgentDraft: @@ -46,7 +48,9 @@ class ClarifyingUserSpecAgent(UserSpecAgent): def __init__(self, runtime: AppRuntime) -> None: super().__init__(runtime) - async def _build_edit_plan(self, user_message: str) -> EditClarificationRequest: + async def _build_edit_plan( + self, user_message: str, conversation_history: list[ConversationMessage] + ) -> EditClarificationRequest: return EditClarificationRequest( question="Which pages should be changed?", reason="The request does not specify the target pages.", diff --git a/frontend/src/prototypes/components/chat/ChatContext.tsx b/frontend/src/prototypes/components/chat/ChatContext.tsx index e6d42e57d7..027bd56959 100644 --- a/frontend/src/prototypes/components/chat/ChatContext.tsx +++ b/frontend/src/prototypes/components/chat/ChatContext.tsx @@ -238,6 +238,8 @@ export function ChatProvider({ children }: { children: ReactNode }) { const { files: activeFiles, fileStubs: activeFileStubs } = useAllFiles(); const { actions: fileActions } = useFileActions(); const abortRef = useRef(null); + const messagesRef = useRef(state.messages); + messagesRef.current = state.messages; // Download a File from the Stirling files endpoint. const downloadFile = useCallback( @@ -323,6 +325,8 @@ export function ChatProvider({ children }: { children: ReactNode }) { const controller = new AbortController(); abortRef.current = controller; + const priorMessages = messagesRef.current; + const userMessage: ChatMessage = { id: crypto.randomUUID(), role: "user", @@ -339,6 +343,10 @@ export function ChatProvider({ children }: { children: ReactNode }) { activeFiles.forEach((file, i) => { formData.append(`fileInputs[${i}].fileInput`, file); }); + priorMessages.forEach((message, i) => { + formData.append(`conversationHistory[${i}].role`, message.role); + formData.append(`conversationHistory[${i}].content`, message.content); + }); const response = await fetch("/api/v1/ai/orchestrate/stream", { method: "POST",