From 2a856fbc19a1030f9dd8c5ebbd72b722cba26ed8 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 000000000..32a811302 --- /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 22228d2aa..da327177c 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 0d02a4438..14bc56ed9 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 e4c17d005..c2537f78c 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 b5ad8c431..40b136578 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 60f7ea94c..c646159da 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 d30b5e53d..40a028aa1 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 21633c075..6fc075b44 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 9ee8856dd..50038d150 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 3fa788fe6..759837288 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 2bcfe7ac6..a7bcedfa9 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 987dc9b59..4ee0d2596 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 1f9de664c..3a91fcf1f 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 e6d42e57d..027bd5695 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",