diff --git a/assistants/prospector-assistant/assistant/agents/document/guided_conversation.py b/assistants/prospector-assistant/assistant/agents/document/guided_conversation.py
index d6ae3250..a3757725 100644
--- a/assistants/prospector-assistant/assistant/agents/document/guided_conversation.py
+++ b/assistants/prospector-assistant/assistant/agents/document/guided_conversation.py
@@ -139,7 +139,10 @@ async def step_conversation(
next_step_name = StepName.DO_DRAFT_OUTLINE
elif user_decision == "draft_paper":
status = Status.USER_COMPLETED
- next_step_name = StepName.DO_FINISH
+ next_step_name = (
+ StepName.DP_DRAFT_CONTENT
+ ) # problem if in draft outline mode... that is supposed to go to DO_FINISH.
+ # coupling is now a problem. and Need to fix the two locations for setting the branching/flow.
else:
logger.error("unknown user decision")
else:
diff --git a/assistants/prospector-assistant/assistant/agents/document/status.py b/assistants/prospector-assistant/assistant/agents/document/status.py
index a5c5e77d..734abc24 100644
--- a/assistants/prospector-assistant/assistant/agents/document/status.py
+++ b/assistants/prospector-assistant/assistant/agents/document/status.py
@@ -15,3 +15,4 @@ class StepName(StrEnum):
DO_DRAFT_OUTLINE = "step_draft_outline"
DO_GC_GET_OUTLINE_FEEDBACK = "step_gc_get_outline_feedback"
DO_FINISH = "step_finish"
+ DP_DRAFT_CONTENT = "step_draft_content"
diff --git a/assistants/prospector-assistant/assistant/agents/document_agent.py b/assistants/prospector-assistant/assistant/agents/document_agent.py
index 2629ecba..9ed800ca 100644
--- a/assistants/prospector-assistant/assistant/agents/document_agent.py
+++ b/assistants/prospector-assistant/assistant/agents/document_agent.py
@@ -39,6 +39,7 @@
class ModeName(StrEnum):
UNDEFINED = "undefined"
DRAFT_OUTLINE = "mode_draft_outline"
+ DRAFT_PAPER = "mode_draft_paper"
class Step(BaseModel):
@@ -224,8 +225,11 @@ class DocumentAgent:
def __init__(self, attachments_extension: AttachmentsExtension) -> None:
self._attachments_extension: AttachmentsExtension = attachments_extension
self._state: State | None = None
- self._commands: list[Callable] = [self._set_mode_draft_outline]
- self._mode_name_to_method: dict[ModeName, Callable] = {ModeName.DRAFT_OUTLINE: self._mode_draft_outline}
+ self._commands: list[Callable] = [self._set_mode_draft_outline, self._set_mode_draft_paper]
+ self._mode_name_to_method: dict[ModeName, Callable] = {
+ ModeName.DRAFT_OUTLINE: self._mode_draft_outline,
+ ModeName.DRAFT_PAPER: self._mode_draft_paper,
+ }
self._step_name_to_method: dict[StepName, Callable] = {} # To be defined in mode method
@property
@@ -302,6 +306,27 @@ def _set_mode_draft_outline(
self._state.mode = Mode(name=ModeName.DRAFT_OUTLINE, status=Status.INITIATED)
self._write_state(context)
+ def _set_mode_draft_paper(
+ self,
+ config: AssistantConfigModel,
+ context: ConversationContext,
+ message: ConversationMessage,
+ metadata: dict[str, Any] = {},
+ ) -> None:
+ # Pre-requisites
+ if self._state is None:
+ logger.error("Document Agent state is None. Returning.")
+ return
+
+ mode = self._state.mode
+ if mode.is_running():
+ logger.warning("Document Agent already in mode: %s. Cannot change modes.", mode.get_name())
+ return
+
+ # Run
+ self._state.mode = Mode(name=ModeName.DRAFT_PAPER, status=Status.INITIATED)
+ self._write_state(context)
+
async def respond_to_conversation(
self,
config: AssistantConfigModel,
@@ -488,6 +513,60 @@ async def _mode_draft_outline(
# Run
return await self._run_mode(config, context, message, metadata)
+ async def _mode_draft_paper(
+ self,
+ config: AssistantConfigModel,
+ context: ConversationContext,
+ message: ConversationMessage,
+ metadata: dict[str, Any] = {},
+ ) -> Status:
+ # Pre-requisites
+ if self._state is None:
+ logger.error("Document Agent state is None. Returning.")
+ return Status.UNDEFINED
+
+ mode = self._state.mode
+ mode_name = mode.get_name()
+ mode_status = mode.get_status()
+
+ if mode_name is not ModeName.DRAFT_PAPER or (
+ mode_status is not Status.NOT_COMPLETED and mode_status is not Status.INITIATED
+ ):
+ logger.error(
+ "Document Agent state mode: %s, mode called: %s, state mode completion status: %s. Resetting Mode.",
+ mode_name,
+ ModeName.DRAFT_PAPER,
+ mode_status,
+ )
+ self._state.mode.reset()
+ self._write_state(context)
+ return self._state.mode.get_status()
+
+ # Setup on first run.
+ if mode_status is Status.INITIATED:
+ self._state.mode.set_step_order(
+ [
+ StepName.DO_GC_ATTACHMENT_CHECK,
+ StepName.DO_DRAFT_OUTLINE,
+ StepName.DO_GC_GET_OUTLINE_FEEDBACK,
+ StepName.DP_DRAFT_CONTENT,
+ ],
+ )
+ logger.info("Document Agent mode (%s) at beginning.", mode_name)
+ first_step_name = self._state.mode.get_step_order()[0]
+ self._state.mode.set_step(Step(name=first_step_name, status=Status.INITIATED))
+ self._write_state(context)
+
+ self._step_name_to_method: dict[StepName, Callable] = {
+ 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.DP_DRAFT_CONTENT: self._step_draft_content,
+ }
+
+ # Run
+ return await self._run_mode(config, context, message, metadata)
+
async def _step_gc_attachment_check(
self,
config: AssistantConfigModel,
@@ -627,6 +706,46 @@ async def _step_finish(
# pretend completed
return Status.USER_COMPLETED, None
+ async def _step_draft_content(
+ self,
+ config: AssistantConfigModel,
+ context: ConversationContext,
+ message: ConversationMessage,
+ metadata: dict[str, Any] = {},
+ ) -> 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, next_step
+
+ step = self._state.mode.get_step()
+ step_name = step.get_name()
+ step_status = step.get_status()
+
+ step_called = StepName.DP_DRAFT_CONTENT
+ if step_name is not step_called or (
+ step_status is not Status.NOT_COMPLETED and step_status is not Status.INITIATED
+ ):
+ logger.error(
+ "Document Agent state step: %s, step called: %s, state step completion status: %s. Resetting Mode.",
+ step_name,
+ step_called,
+ step_status,
+ )
+ self._state.mode.reset()
+ self._write_state(context)
+ return self._state.mode.get_status(), next_step
+
+ # Run
+ logger.info("Document Agent running step: %s", step_name)
+ status, next_step_name = await self._draft_content(config, context, message, metadata)
+ step.set_status(status)
+ self._state.mode.set_step(step)
+ self._write_state(context)
+ return step.get_status(), next_step_name
+
# endregion
#
@@ -726,7 +845,7 @@ async def _draft_outline(
# create chat completion messages
chat_completion_messages: list[ChatCompletionMessageParam] = []
- chat_completion_messages.append(_main_system_message())
+ chat_completion_messages.append(_draft_outline_main_system_message())
chat_completion_messages.append(
_chat_history_system_message(conversation.messages, participants_list.participants)
)
@@ -743,19 +862,19 @@ async def _draft_outline(
"response_format": {"type": "text"},
}
completion = await client.chat.completions.create(**completion_args)
- content = completion.choices[0].message.content
+ message_content = completion.choices[0].message.content
_on_success_metadata_update(metadata, method_metadata_key, config, chat_completion_messages, completion)
except Exception as e:
logger.exception(f"exception occurred calling openai chat completion: {e}")
- content = (
+ message_content = (
"An error occurred while calling the OpenAI API. Is it configured correctly?"
"View the debug inspector for more information."
)
_on_error_metadata_update(metadata, method_metadata_key, config, chat_completion_messages, e)
# store only latest version for now (will keep all versions later as need arises)
- (storage_directory_for_context(context) / "document_agent/outline.txt").write_text(content)
+ (storage_directory_for_context(context) / "document_agent/outline.txt").write_text(message_content)
# send the response to the conversation only if from a command. Otherwise return info to caller.
message_type = MessageType.chat
@@ -764,7 +883,7 @@ async def _draft_outline(
await context.send_messages(
NewConversationMessage(
- content=content,
+ content=message_content,
message_type=message_type,
metadata=metadata,
)
@@ -851,6 +970,85 @@ async def _gc_outline_feedback(
return conversation_status, next_step_name
+ async def _draft_content(
+ self,
+ config: AssistantConfigModel,
+ context: ConversationContext,
+ message: ConversationMessage,
+ metadata: dict[str, Any] = {},
+ ) -> tuple[Status, StepName | None]:
+ method_metadata_key = "draft_content"
+
+ # get conversation related info
+ conversation = await context.get_messages(before=message.id)
+ if message.message_type == MessageType.chat:
+ conversation.messages.append(message)
+ participants_list = await context.get_participants(include_inactive=True)
+
+ # get attachments related info
+ attachment_messages = await self._attachments_extension.get_completion_messages_for_attachments(
+ context, config=config.agents_config.attachment_agent
+ )
+
+ # get outline related info
+ outline: str | None = None
+ content: str | None = None
+ # path = _get_document_agent_conversation_storage_path(context)
+ if path.exists(storage_directory_for_context(context) / "document_agent/outline.txt"):
+ outline = (storage_directory_for_context(context) / "document_agent/outline.txt").read_text()
+ if path.exists(storage_directory_for_context(context) / "document_agent/content.txt"):
+ content = (storage_directory_for_context(context) / "document_agent/content.txt").read_text()
+
+ # create chat completion messages
+ chat_completion_messages: list[ChatCompletionMessageParam] = []
+ chat_completion_messages.append(_draft_content_main_system_message())
+ chat_completion_messages.append(
+ _chat_history_system_message(conversation.messages, participants_list.participants)
+ )
+ chat_completion_messages.extend(attachment_messages)
+ if outline is not None:
+ chat_completion_messages.append(_outline_system_message(outline))
+ if content is not None: # only grabs previously written content, not all yet.
+ chat_completion_messages.append(_content_system_message(content))
+
+ # make completion call to openai
+ async with openai_client.create_client(config.service_config) as client:
+ try:
+ completion_args = {
+ "messages": chat_completion_messages,
+ "model": config.request_config.openai_model,
+ "response_format": {"type": "text"},
+ }
+ completion = await client.chat.completions.create(**completion_args)
+ content = completion.choices[0].message.content
+ _on_success_metadata_update(metadata, method_metadata_key, config, chat_completion_messages, completion)
+
+ except Exception as e:
+ logger.exception(f"exception occurred calling openai chat completion: {e}")
+ content = (
+ "An error occurred while calling the OpenAI API. Is it configured correctly?"
+ "View the debug inspector for more information."
+ )
+ _on_error_metadata_update(metadata, method_metadata_key, config, chat_completion_messages, e)
+
+ # store only latest version for now (will keep all versions later as need arises)
+ (storage_directory_for_context(context) / "document_agent/content.txt").write_text(content)
+
+ # send the response to the conversation only if from a command. Otherwise return info to caller.
+ message_type = MessageType.chat
+ if message.message_type == MessageType.command:
+ message_type = MessageType.command
+
+ await context.send_messages(
+ NewConversationMessage(
+ content=content,
+ message_type=message_type,
+ metadata=metadata,
+ )
+ )
+
+ return Status.USER_COMPLETED, None
+
# endregion
@@ -859,11 +1057,19 @@ async def _gc_outline_feedback(
#
-def _main_system_message() -> ChatCompletionSystemMessageParam:
+def _draft_outline_main_system_message() -> ChatCompletionSystemMessageParam:
message: ChatCompletionSystemMessageParam = {"role": "system", "content": draft_outline_main_system_message}
return message
+def _draft_content_main_system_message() -> ChatCompletionSystemMessageParam:
+ message: ChatCompletionSystemMessageParam = {
+ "role": "system",
+ "content": draft_content_continue_main_system_message,
+ }
+ return message
+
+
def _chat_history_system_message(
conversation_messages: list[ConversationMessage],
participants: list[ConversationParticipant],
@@ -882,11 +1088,18 @@ def _chat_history_system_message(
def _outline_system_message(outline: str) -> ChatCompletionSystemMessageParam:
- if outline is not None:
- message: ChatCompletionSystemMessageParam = {
- "role": "system",
- "content": (f"{outline}"),
- }
+ message: ChatCompletionSystemMessageParam = {
+ "role": "system",
+ "content": (f"{outline}"),
+ }
+ return message
+
+
+def _content_system_message(content: str) -> ChatCompletionSystemMessageParam:
+ message: ChatCompletionSystemMessageParam = {
+ "role": "system",
+ "content": (f"{content}"),
+ }
return message
@@ -904,6 +1117,29 @@ def _outline_system_message(outline: str) -> ChatCompletionSystemMessageParam:
# " outline if it exists, as well as any other attachments provided by the user to inform a newly revised "
# "outline draft. Provide ONLY any outline. Provide no further instructions to the user.")
+draft_content_continue_main_system_message = (
+ "Following the structure of the outline, create the content for the next (or first) page of the"
+ " document - don't try to create the entire document in one pass nor wrap it up too quickly, it will be a"
+ " multi-page document so just create the next page. It's more important to maintain"
+ " an appropriately useful level of detail. After this page is generated, the system will follow up"
+ " and ask for the next page. If you have already generated all the pages for the"
+ " document as defined in the outline, return empty content."
+)
+# ("You are an AI assistant that helps draft new content of a document based on an outline."
+# " You use information from a chat history between a user and an assistant, the approved outline from the user,"
+# "and an existing version of drafted content if it exists, as well as any other attachments provided by the user to inform newly revised "
+# "content. Newly drafted content does not need to cover the entire outline. Instead it should be limited to a reasonable 100 lines of natural language"
+# " or subsection of the outline (which ever is shorter). The newly drafted content should be written as to append to any existing drafted content."
+# " This way the user can review newly drafted content as a subset of the future full document and not be overwhelmed."
+# "Only provide the newly drafted content. Provide no further instructions to the user.")
+
+draft_content_iterate_main_system_message = (
+ "Following the structure of the outline, iterate on the currently drafted page of the"
+ " document. It's more important to maintain"
+ " an appropriately useful level of detail. After this page is iterated upon, the system will follow up"
+ " and ask for the next page."
+)
+
def _on_success_metadata_update(
metadata: dict[str, Any],
@@ -969,19 +1205,3 @@ def _format_message(message: ConversationMessage, participants: list[Conversatio
# endregion
-
-
-##### FROM NOTEBOOK
-# await document_skill.draft_outline(context=unused, openai_client=async_client, model=model)
-#
-# decision, user_feedback = await document_skill.get_user_feedback(
-# context=unused, openai_client=async_client, model=model, outline=True
-# )
-#
-# while decision == "[ITERATE]":
-# await document_skill.draft_outline(
-# context=unused, openai_client=async_client, model=model, user_feedback=user_feedback
-# )
-# decision, user_feedback = await document_skill.get_user_feedback(
-# context=unused, openai_client=async_client, model=model, outline=True
-# )