Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix routing to next step per user input #225

Merged
merged 13 commits into from
Nov 7, 2024
Merged
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"
109 changes: 54 additions & 55 deletions assistants/prospector-assistant/assistant/agents/document_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand All @@ -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"
Expand Down Expand Up @@ -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:
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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()
Expand All @@ -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()
Expand All @@ -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()
Expand All @@ -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
Expand All @@ -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

Expand All @@ -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()
Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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

Expand Down