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",