Fix prospector. (microsoft#132)

Update attachment agent to attachment extension in document agent.

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

@@ -3,21 +3,22 @@
 
 import deepmerge
 import openai_client
-from openai.types.chat import ChatCompletionMessageParam, ChatCompletionSystemMessageParam
+from assistant_extensions.attachments import AttachmentsExtension
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionSystemMessageParam,
+)
 from semantic_workbench_api_model.workbench_model import (
-    #    AssistantStateEvent,
     ConversationMessage,
     ConversationParticipant,
     MessageType,
     NewConversationMessage,
-    #    ParticipantRole,
 )
 from semantic_workbench_assistant.assistant_app import (
     ConversationContext,
     #    storage_directory_for_context,
 )
 
-from ..agents.attachment_agent import AttachmentAgent
 from ..config import AssistantConfigModel
 
 logger = logging.getLogger(__name__)
@@ -32,7 +33,8 @@ class DocumentAgent:
     An agent for working on document content: creation, editing, translation, etc.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, attachments_extension: AttachmentsExtension) -> None:
+        self.attachments_extension = attachments_extension
         self._commands = [self.draft_outline]
 
     @property
@@ -76,15 +78,18 @@ async def draft_outline(
             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
+        )
+
         # create chat completion messages
         chat_completion_messages: list[ChatCompletionMessageParam] = []
         _add_main_system_message(chat_completion_messages, draft_outline_main_system_message)
         _add_chat_history_system_message(
             chat_completion_messages, conversation.messages, participants_list.participants
         )
-        _add_attachments_system_message(
-            chat_completion_messages, config, AttachmentAgent.generate_attachment_messages(context)
-        )
+        chat_completion_messages.extend(attachment_messages)
 
         # make completion call to openai
         async with openai_client.create_client(config.service_config) as client:
@@ -150,19 +155,6 @@ def _add_chat_history_system_message(
     chat_completion_messages.append(message)
 
 
-def _add_attachments_system_message(
-    chat_completion_messages: list[ChatCompletionMessageParam],
-    config: AssistantConfigModel,
-    attachment_messages: list[ChatCompletionMessageParam],
-) -> None:
-    if len(attachment_messages) > 0:
-        chat_completion_messages.append({
-            "role": "system",
-            "content": config.agents_config.attachment_agent.context_description,
-        })
-        chat_completion_messages.extend(attachment_messages)
-
-
 draft_outline_main_system_message = (
     "Generate an outline for the document, including title. The outline should include the key points that will"
     " be covered in the document. If attachments exist, consider the attachments and the rationale for why they"
@@ -199,7 +191,6 @@ def _on_success_metadata_update(
             }
         },
     )
-    metadata = _reduce_metadata_debug_payload(metadata)
 
 
 def _on_error_metadata_update(
@@ -223,7 +214,6 @@ def _on_error_metadata_update(
             }
         },
     )
-    metadata = _reduce_metadata_debug_payload(metadata)
 
 
 # endregion
@@ -247,58 +237,4 @@ def _format_message(message: ConversationMessage, participants: list[Conversatio
     return f"[{participant_name} - {message_datetime}]: {message.content}"
 
 
-def _reduce_metadata_debug_payload(metadata: dict[str, Any]) -> dict[str, Any]:
-    """
-    Reduce the size of the metadata debug payload.
-    """
-
-    # map of payload reducers and keys they should be called for
-    # each reducer should take a value and return a reduced value
-    payload_reducers: dict[str, list[Any]] = {
-        "content": [
-            AttachmentAgent.reduce_attachment_payload_from_content,
-        ]
-    }
-
-    # NOTE: use try statements around each recursive call to reduce_metadata to report the parent key in case of error
-
-    # now iterate recursively over all metadata keys and call the payload reducers for the matching keys
-    def reduce_metadata(metadata: dict[str, Any] | Any) -> dict[str, Any] | Any:
-        # check if the metadata is not a dictionary
-        if not isinstance(metadata, dict):
-            return metadata
-
-        # iterate over the metadata keys
-        for key, value in metadata.items():
-            # check if the key is in the payload reducers
-            if key in payload_reducers:
-                # call the payload reducer for the key
-                for reducer in payload_reducers[key]:
-                    metadata[key] = reducer(value)
-            # check if the value is a dictionary
-            if isinstance(value, dict):
-                try:
-                    # recursively reduce the metadata
-                    metadata[key] = reduce_metadata(value)
-                except Exception as e:
-                    logger.exception(f"exception occurred reducing metadata for key '{key}': {e}")
-            # check if the value is a list
-            elif isinstance(value, list):
-                try:
-                    # recursively reduce the metadata for each item in the list
-                    metadata[key] = [reduce_metadata(item) for item in value]
-                except Exception as e:
-                    logger.exception(f"exception occurred reducing metadata for key '{key}': {e}")
-
-        # return the reduced metadata
-        return metadata
-
-    try:
-        # reduce the metadata
-        return reduce_metadata(metadata)
-    except Exception as e:
-        logger.exception(f"exception occurred reducing metadata: {e}")
-        return metadata
-
-
 # endregion

assistants/prospector-assistant/assistant/chat.py

@@ -112,17 +112,19 @@ async def content_evaluator_factory(context: ConversationContext) -> ContentSafe
 async def on_command_message_created(
     context: ConversationContext, event: ConversationEvent, message: ConversationMessage
 ) -> None:
+    config = await assistant_config.get(context.assistant)
+    metadata: dict[str, Any] = {"debug": {"content_safety": event.data.get(content_safety.metadata_key, {})}}
+
     # For now, handling only commands from Document Agent for exploration of implementation
     # We assume Document Agent is available and future logic would determine which agent
     # the command is intended for. Assumption made in order to make doc agent available asap.
-    metadata: dict[str, Any] = {"debug": {"content_safety": event.data.get(content_safety.metadata_key, {})}}
 
-    config = await assistant_config.get(context.assistant)
-    if config.agents_config.attachment_agent.include_in_response_generation:
-        doc_agent = DocumentAgent()
-        await doc_agent.receive_command(config, context, message, metadata)
-    else:
-        pass  # for now
+    # We should not be creating agent instances for commands and conversation separately.
+    # This will need to be resolved.
+
+    # if config.agents_config.document_agent.enabled:
+    doc_agent = DocumentAgent(attachments_extension)
+    await doc_agent.receive_command(config, context, message, metadata)
 
 
 @assistant.events.conversation.message.chat.on_created