Fix routing to next step per user input (#225)

- Adds new routing to a different step if gc (or any step function)
returns a next step.
- Updates gc config for outline feedback to attempt to fix issues.
(ongoing)

assistants/prospector-assistant/assistant/agents/document/gc_draft_outline_feedback_config.py

@@ -24,24 +24,29 @@ class ArtifactModel(BaseModel):
 
 # Rules - These are the do's and don'ts that the agent should follow during the conversation.
 rules = [
+    "Do NOT rewrite or update the outline, even if the user asks you to."
+    "You are ONLY allowed to help the user decide on any changes to the outline or answer questions about writing an outline."
     "Terminate the conversation immediately if the user asks for harmful or inappropriate content.",
+    "If the conversation_status is marked as user_complete, the final_response cannot be left as 'Unanswered'. The final_response must be set based on the conversation flow instructions.",
 ]
 
 # Conversation Flow (optional) - This defines in natural language the steps of the conversation.
 conversation_flow = """1. Start by asking the user to review the drafted outline.
 2. Answer any questions about the outline or the drafting process the user might want to explore.
 3. At any time, if the user asks for a change to the outline or updates the attachment file list, the conversation_status must be
-marked as user_completed. The user_decision must be marked as update_outline. The final_response should inform the user that a new outline is being generated based off the new info or request.
+marked as user_completed. The user_decision must be marked as update_outline. The final_response must inform the user that a new outline is being generated based off the new info or request.
 4. At any time, if the user is good with the outline in its current form and ready to move on to drafting a paper from it, the conversation_status must be
-marked as user_completed. The user_decision must be marked as draft_paper. The final_response should inform the user that you will start drafting the beginning of the
+marked as user_completed. The user_decision must be marked as draft_paper. The final_response must inform the user that you will start drafting the beginning of the
 document based on this outline.
 """
 
 # Context (optional) - This is any additional information or the circumstances the agent is in that it should be aware of.
 # It can also include the high level goal of the conversation if needed.
 context = """You are working with a user on drafting an outline. The current drafted outline is provided, along with any filenames
 that were used to help draft the outline. You do not have access to the content within the filenames that were used to help draft the outline.
- Your purpose here is to help the user decide on any changes to the outline they might want or answer questions about it."""
+ Your purpose here is to help the user decide on any changes to the outline they might want or answer questions about it.  This may be the first time
+ the user is asking for you help, or the nth time.  Please use the conversation history provided to determine if you should be give an initial greeting
+ to the user or continuing the draft outline process."""
 
 
 # Resource Constraints (optional) - This defines the constraints on the conversation such as time or turns.

assistants/prospector-assistant/assistant/agents/document/guided_conversation.py

@@ -13,7 +13,7 @@
 
 from ...config import AssistantConfigModel
 from .config import GuidedConversationAgentConfigModel
-from .status import Status
+from .status import Status, StepName
 
 logger = logging.getLogger(__name__)
 
@@ -44,10 +44,11 @@ async def step_conversation(
         agent_config: GuidedConversationAgentConfigModel,
         conversation_context: ConversationContext,
         last_user_message: str | None,
-    ) -> tuple[str, Status]:
+    ) -> tuple[str, Status, StepName | None]:
         """
         Step the conversation to the next turn.
         """
+        next_step_name = None
 
         rules = agent_config.rules
         conversation_flow = agent_config.conversation_flow
@@ -134,17 +135,19 @@ async def step_conversation(
                     status = Status.USER_COMPLETED
                 else:
                     if user_decision == "update_outline":  # this code is becoming highly coupled fyi to the gc configs
-                        status = Status.UPDATE_OUTLINE
+                        status = Status.USER_COMPLETED
+                        next_step_name = StepName.DO_DRAFT_OUTLINE
                     elif user_decision == "draft_paper":
                         status = Status.USER_COMPLETED
+                        next_step_name = StepName.DO_FINISH
                     else:
                         logger.error("unknown user decision")
             else:
                 _delete_guided_conversation_state(conversation_context)
                 status = Status.USER_EXIT_EARLY
                 response = final_response
 
-        return response, status
+        return response, status, next_step_name
 
     # endregion
 

assistants/prospector-assistant/assistant/agents/document/status.py

@@ -5,6 +5,13 @@ class Status(StrEnum):
     UNDEFINED = "undefined"
     INITIATED = "initiated"
     NOT_COMPLETED = "not_completed"
-    UPDATE_OUTLINE = "update_outline"
     USER_COMPLETED = "user_completed"
     USER_EXIT_EARLY = "user_exit_early"
+
+
+class StepName(StrEnum):
+    UNDEFINED = "undefined"
+    DO_GC_ATTACHMENT_CHECK = "step_gc_attachment_check"
+    DO_DRAFT_OUTLINE = "step_draft_outline"
+    DO_GC_GET_OUTLINE_FEEDBACK = "step_gc_get_outline_feedback"
+    DO_FINISH = "step_finish"

assistants/prospector-assistant/assistant/agents/document_agent.py

@@ -26,7 +26,7 @@
 from .document.gc_attachment_check_config import GCAttachmentCheckConfigModel
 from .document.gc_draft_outline_feedback_config import GCDraftOutlineFeedbackConfigModel
 from .document.guided_conversation import GuidedConversationAgent
-from .document.status import Status
+from .document.status import Status, StepName
 
 logger = logging.getLogger(__name__)
 
@@ -36,14 +36,6 @@
 #
 
 
-class StepName(StrEnum):
-    UNDEFINED = "undefined"
-    DO_GC_ATTACHMENT_CHECK = "step_gc_attachment_check"
-    DO_DRAFT_OUTLINE = "step_draft_outline"
-    DO_GC_GET_OUTLINE_FEEDBACK = "step_gc_get_outline_feedback"
-    DO_FINAL_OUTLINE = "step_final_outline"
-
-
 class ModeName(StrEnum):
     UNDEFINED = "undefined"
     DRAFT_OUTLINE = "mode_draft_outline"
@@ -385,7 +377,7 @@ async def _run_mode(
             step_method = self._get_step_method(step)
             if step_method:
                 logger.info("Document Agent in step: %s", step_name)
-                step_status = await step_method(config, context, message, metadata)
+                step_status, next_step_name = await step_method(config, context, message, metadata)
 
                 match step_status:  # resulting status of step_method()
                     case Status.NOT_COMPLETED:
@@ -394,18 +386,26 @@ async def _run_mode(
                         break  # ok - get more user input
 
                     case Status.USER_COMPLETED:
-                        next_step = self._state.mode.get_next_step()
-                        if next_step is not None:
-                            step = next_step
-                            step_name = next_step.get_name()
-                            step_status = next_step.get_status()  # new step is Status.INITIATED
-                            self._state.mode.set_step(next_step)
+                        if next_step_name is not None:
+                            step = Step(name=next_step_name, status=Status.INITIATED)
+                            step_name = step.get_name()
+                            step_status = step.get_status()
+                            self._state.mode.set_step(step)
                             self._write_state(context)
                             continue  # ok - don't need user input yet
-                        else:
-                            self._state.mode.get_step().set_status(step_status)
-                            self._state.mode.set_status(step_status)
-                            break  # ok - all done :)
+                        else:  # go with prescribed order?  Not sure if we want this long term, or just go with above
+                            next_step = self._state.mode.get_next_step()
+                            if next_step is not None:
+                                step = next_step
+                                step_name = next_step.get_name()
+                                step_status = next_step.get_status()  # new step is Status.INITIATED
+                                self._state.mode.set_step(next_step)
+                                self._write_state(context)
+                                continue  # ok - don't need user input yet
+                            else:
+                                self._state.mode.get_step().set_status(step_status)
+                                self._state.mode.set_status(step_status)
+                                break  # ok - all done :)
 
                     case Status.USER_EXIT_EARLY:
                         self._state.mode.get_step().set_status(step_status)
@@ -470,7 +470,7 @@ async def _mode_draft_outline(
                     StepName.DO_GC_ATTACHMENT_CHECK,
                     StepName.DO_DRAFT_OUTLINE,
                     StepName.DO_GC_GET_OUTLINE_FEEDBACK,
-                    StepName.DO_FINAL_OUTLINE,
+                    StepName.DO_FINISH,
                 ],
             )
             logger.info("Document Agent mode (%s) at beginning.", mode_name)
@@ -482,7 +482,7 @@ async def _mode_draft_outline(
             StepName.DO_GC_ATTACHMENT_CHECK: self._step_gc_attachment_check,
             StepName.DO_DRAFT_OUTLINE: self._step_draft_outline,
             StepName.DO_GC_GET_OUTLINE_FEEDBACK: self._step_gc_get_outline_feedback,
-            StepName.DO_FINAL_OUTLINE: self._step_final_outline,
+            StepName.DO_FINISH: self._step_finish,
         }
 
         # Run
@@ -494,11 +494,13 @@ async def _step_gc_attachment_check(
         context: ConversationContext,
         message: ConversationMessage,
         metadata: dict[str, Any] = {},
-    ) -> Status:
+    ) -> tuple[Status, StepName | None]:
+        next_step = None
+
         # Pre-requisites
         if self._state is None:
             logger.error("Document Agent state is None. Returning.")
-            return Status.UNDEFINED
+            return Status.UNDEFINED, next_step
 
         step = self._state.mode.get_step()
         step_name = step.get_name()
@@ -517,27 +519,29 @@ async def _step_gc_attachment_check(
             )
             self._state.mode.reset()
             self._write_state(context)
-            return self._state.mode.get_status()
+            return self._state.mode.get_status(), next_step
 
         # Run
         logger.info("Document Agent running step: %s", step_name)
-        status = await self._gc_attachment_check(config, context, message, metadata)
+        status, next_step_name = await self._gc_attachment_check(config, context, message, metadata)
         step.set_status(status)
         self._state.mode.set_step(step)
         self._write_state(context)
-        return step.get_status()
+        return step.get_status(), next_step_name
 
     async def _step_draft_outline(
         self,
         config: AssistantConfigModel,
         context: ConversationContext,
         message: ConversationMessage,
         metadata: dict[str, Any] = {},
-    ) -> Status:
+    ) -> tuple[Status, StepName | None]:
+        next_step = None
+
         # Pre-requisites
         if self._state is None:
             logger.error("Document Agent state is None. Returning.")
-            return Status.UNDEFINED
+            return Status.UNDEFINED, next_step
 
         step = self._state.mode.get_step()
         step_name = step.get_name()
@@ -555,27 +559,29 @@ async def _step_draft_outline(
             )
             self._state.mode.reset()
             self._write_state(context)
-            return self._state.mode.get_status()
+            return self._state.mode.get_status(), next_step
 
         # Run
         logger.info("Document Agent running step: %s", step_name)
-        status = await self._draft_outline(config, context, message, metadata)
+        status, next_step_name = await self._draft_outline(config, context, message, metadata)
         step.set_status(status)
         self._state.mode.set_step(step)
         self._write_state(context)
-        return step.get_status()
+        return step.get_status(), next_step_name
 
     async def _step_gc_get_outline_feedback(
         self,
         config: AssistantConfigModel,
         context: ConversationContext,
         message: ConversationMessage,
         metadata: dict[str, Any] = {},
-    ) -> Status:
+    ) -> tuple[Status, StepName | None]:
+        next_step_name = None
+
         # Pre-requisites
         if self._state is None:
             logger.error("Document Agent state is None. Returning.")
-            return Status.UNDEFINED
+            return Status.UNDEFINED, next_step_name
 
         step = self._state.mode.get_step()
         step_name = step.get_name()
@@ -594,7 +600,7 @@ async def _step_gc_get_outline_feedback(
             )
             self._state.mode.reset()
             self._write_state(context)
-            return self._state.mode.get_status()
+            return self._state.mode.get_status(), next_step_name
 
         # Run
         user_message: ConversationMessage | None
@@ -604,29 +610,22 @@ async def _step_gc_get_outline_feedback(
             user_message = message
 
         logger.info("Document Agent running step: %s", step_name)
-        status = await self._gc_outline_feedback(config, context, user_message, metadata)
-
-        # should this be status or a different return as part of a tuple... like next step request?
-        # this way the control would still be higher (and non recursive) as the branching logic will return instead of call internally the next step.
-        # but controlling of what gets called next will be determined by each step at its completion.  it seems like there should still be
-        # some level of oversight by the mode layer in case we end up in an endless loop of steps saying the should each be called next...(still a problem.)
-        if status is Status.UPDATE_OUTLINE:
-            status = await self._draft_outline(config, context, message, metadata)
+        status, next_step_name = await self._gc_outline_feedback(config, context, user_message, metadata)
 
         step.set_status(status)
         self._state.mode.set_step(step)
         self._write_state(context)
-        return step.get_status()
+        return step.get_status(), next_step_name
 
-    async def _step_final_outline(
+    async def _step_finish(
         self,
         config: AssistantConfigModel,
         context: ConversationContext,
         message: ConversationMessage,
         metadata: dict[str, Any] = {},
-    ) -> Status:
+    ) -> tuple[Status, StepName | None]:
         # pretend completed
-        return Status.USER_COMPLETED
+        return Status.USER_COMPLETED, None
 
     # endregion
 
@@ -640,7 +639,7 @@ async def _gc_attachment_check(
         context: ConversationContext,
         message: ConversationMessage,
         metadata: dict[str, Any] = {},
-    ) -> Status:
+    ) -> tuple[Status, StepName | None]:
         method_metadata_key = "document_agent_gc_response"
 
         gc_convo_config: GuidedConversationAgentConfigModel = GCAttachmentCheckConfigModel()
@@ -654,7 +653,7 @@ async def _gc_attachment_check(
         gc_convo_config.context = gc_convo_config.context + "\n\n" + filenames_str
 
         try:
-            response_message, conversation_status = await GuidedConversationAgent.step_conversation(
+            response_message, conversation_status, next_step_name = await GuidedConversationAgent.step_conversation(
                 config=config,
                 openai_client=openai_client.create_client(config.service_config),
                 agent_config=gc_convo_config,
@@ -697,15 +696,15 @@ async def _gc_attachment_check(
         # Need to add a good way to stop mode if an exception occurs.
         # Also need to update the gc state turn count to 0 (and any thing else that needs to be reset) once conversation is over... or exception occurs?)
 
-        return conversation_status
+        return conversation_status, next_step_name
 
     async def _draft_outline(
         self,
         config: AssistantConfigModel,
         context: ConversationContext,
         message: ConversationMessage,
         metadata: dict[str, Any] = {},
-    ) -> Status:
+    ) -> tuple[Status, StepName | None]:
         method_metadata_key = "draft_outline"
 
         # get conversation related info
@@ -771,15 +770,15 @@ async def _draft_outline(
             )
         )
 
-        return Status.USER_COMPLETED
+        return Status.USER_COMPLETED, None
 
     async def _gc_outline_feedback(
         self,
         config: AssistantConfigModel,
         context: ConversationContext,
         message: ConversationMessage | None,
         metadata: dict[str, Any] = {},
-    ) -> Status:
+    ) -> tuple[Status, StepName | None]:
         method_metadata_key = "document_agent_gc_response"
 
         gc_do_feedback_config: GuidedConversationAgentConfigModel = GCDraftOutlineFeedbackConfigModel()
@@ -807,7 +806,7 @@ async def _gc_outline_feedback(
             user_message = None
 
         try:
-            response_message, conversation_status = await GuidedConversationAgent.step_conversation(
+            response_message, conversation_status, next_step_name = await GuidedConversationAgent.step_conversation(
                 config=config,
                 openai_client=openai_client.create_client(config.service_config),
                 agent_config=gc_do_feedback_config,
@@ -850,7 +849,7 @@ async def _gc_outline_feedback(
         # Need to add a good way to stop mode if an exception occurs.
         # Also need to update the gc state turn count to 0 (and any thing else that needs to be reset) once conversation is over... or exception occurs?)
 
-        return conversation_status
+        return conversation_status, next_step_name
 
     # endregion