diff --git a/assistants/codespace-assistant/assistant/response/request_builder.py b/assistants/codespace-assistant/assistant/response/request_builder.py index 077ceddd..4e1693ce 100644 --- a/assistants/codespace-assistant/assistant/response/request_builder.py +++ b/assistants/codespace-assistant/assistant/response/request_builder.py @@ -148,7 +148,7 @@ async def build_request( model=request_config.model, ) - logging.info( + logger.info( "chat message params budgeted; message count: %d, total token count: %d", len(chat_message_params), total_token_count, diff --git a/assistants/codespace-assistant/assistant/response/response.py b/assistants/codespace-assistant/assistant/response/response.py index 71a48f9f..8a0857f2 100644 --- a/assistants/codespace-assistant/assistant/response/response.py +++ b/assistants/codespace-assistant/assistant/response/response.py @@ -2,7 +2,14 @@ from contextlib import AsyncExitStack from typing import Any -from assistant_extensions.chat_context_toolkit.archive import ArchiveTaskQueues +from assistant_extensions.attachments import get_attachments +from assistant_extensions.chat_context_toolkit.archive import ( + ArchiveTaskQueues, + construct_archive_summarizer, +) +from assistant_extensions.chat_context_toolkit.message_history import ( + construct_attachment_summarizer, +) from assistant_extensions.mcp import ( MCPClientSettings, MCPServerConnectionError, @@ -166,8 +173,19 @@ async def message_handler(message) -> None: # enqueue an archive task for this conversation await archive_task_queues.enqueue_run( context=context, - service_config=service_config, - request_config=request_config, + attachments=list( + await get_attachments( + context, + summarizer=construct_attachment_summarizer( + service_config=service_config, + request_config=request_config, + ), + ) + ), + archive_summarizer=construct_archive_summarizer( + service_config=service_config, + request_config=request_config, + ), archive_task_config=ArchiveTaskConfig( chunk_token_count_threshold=config.chat_context_config.archive_token_threshold ), diff --git a/assistants/codespace-assistant/assistant/response/step_handler.py b/assistants/codespace-assistant/assistant/response/step_handler.py index c919d45e..a64cff75 100644 --- a/assistants/codespace-assistant/assistant/response/step_handler.py +++ b/assistants/codespace-assistant/assistant/response/step_handler.py @@ -4,7 +4,11 @@ from typing import Any, List import deepmerge -from assistant_extensions.chat_context_toolkit.message_history import chat_context_toolkit_message_provider_for +from assistant_extensions.attachments import get_attachments +from assistant_extensions.chat_context_toolkit.message_history import ( + chat_context_toolkit_message_provider_for, + construct_attachment_summarizer, +) from assistant_extensions.chat_context_toolkit.virtual_filesystem import ( archive_file_source_mount, attachments_file_source_mount, @@ -100,8 +104,15 @@ async def handle_error(error_message: str, error_debug: dict[str, Any] | None = history_message_provider = chat_context_toolkit_message_provider_for( context=context, tool_abbreviations=abbreviations.tool_abbreviations, - service_config=service_config, - request_config=request_config, + attachments=list( + await get_attachments( + context, + summarizer=construct_attachment_summarizer( + service_config=service_config, + request_config=request_config, + ), + ) + ), ) build_request_result = await build_request( diff --git a/assistants/document-assistant/assistant/response/responder.py b/assistants/document-assistant/assistant/response/responder.py index dc43f068..7653b1a9 100644 --- a/assistants/document-assistant/assistant/response/responder.py +++ b/assistants/document-assistant/assistant/response/responder.py @@ -9,7 +9,11 @@ import deepmerge import pendulum -from assistant_extensions.chat_context_toolkit.message_history import chat_context_toolkit_message_provider_for +from assistant_extensions.attachments import get_attachments +from assistant_extensions.chat_context_toolkit.message_history import ( + chat_context_toolkit_message_provider_for, + construct_attachment_summarizer, +) from assistant_extensions.chat_context_toolkit.virtual_filesystem import ( archive_file_source_mount, ) @@ -400,9 +404,15 @@ async def _construct_prompt(self) -> tuple[list, list[ChatCompletionMessageParam message_provider = chat_context_toolkit_message_provider_for( context=self.context, tool_abbreviations=tool_abbreviations, - # use the fast client config for the attachment summarization that the message provider does - service_config=self.config.generative_ai_fast_client_config.service_config, - request_config=self.config.generative_ai_fast_client_config.request_config, + attachments=list( + await get_attachments( + self.context, + summarizer=construct_attachment_summarizer( + service_config=self.config.generative_ai_fast_client_config.service_config, + request_config=self.config.generative_ai_fast_client_config.request_config, + ), + ) + ), ) system_prompt_token_count = num_tokens_from_message(main_system_prompt, model="gpt-4o") tool_token_count = num_tokens_from_tools(tools, model="gpt-4o") @@ -421,7 +431,7 @@ async def _construct_prompt(self) -> tuple[list, list[ChatCompletionMessageParam chat_history: list[ChatCompletionMessageParam] = list(budgeted_messages_result.messages) chat_history.insert(0, main_system_prompt) - logging.info("The system prompt has been constructed.") + logger.info("The system prompt has been constructed.") # Update telemetry for inspector self.latest_telemetry.system_prompt_tokens = system_prompt_token_count self.latest_telemetry.tool_tokens = tool_token_count diff --git a/libraries/python/assistant-extensions/assistant_extensions/attachments/_attachments.py b/libraries/python/assistant-extensions/assistant_extensions/attachments/_attachments.py index f67e9bf4..34fda783 100644 --- a/libraries/python/assistant-extensions/assistant_extensions/attachments/_attachments.py +++ b/libraries/python/assistant-extensions/assistant_extensions/attachments/_attachments.py @@ -118,7 +118,7 @@ async def get_completion_messages_for_attachments( self, context: ConversationContext, config: AttachmentsConfigModel, - include_filenames: list[str] | None = None, + include_filenames: list[str] = [], exclude_filenames: list[str] = [], summarizer: Summarizer | None = None, ) -> Sequence[CompletionMessage]: @@ -143,6 +143,7 @@ async def get_completion_messages_for_attachments( error_handler=self._error_handler, include_filenames=include_filenames, exclude_filenames=exclude_filenames, + summarizer=summarizer, ) if not attachments: @@ -159,14 +160,14 @@ async def get_completion_messages_for_attachments( async def get_attachment_filenames( self, context: ConversationContext, - include_filenames: list[str] | None = None, + include_filenames: list[str] = [], exclude_filenames: list[str] = [], ) -> list[str]: files_response = await context.list_files() # for all files, get the attachment for file in files_response.files: - if include_filenames is not None and file.filename not in include_filenames: + if include_filenames and file.filename not in include_filenames: continue if file.filename in exclude_filenames: continue @@ -226,11 +227,11 @@ async def default_error_handler(context: ConversationContext, filename: str, e: async def get_attachments( context: ConversationContext, - include_filenames: list[str] | None, - exclude_filenames: list[str], + exclude_filenames: list[str] = [], + include_filenames: list[str] = [], error_handler: AttachmentProcessingErrorHandler = default_error_handler, summarizer: Summarizer | None = None, -) -> Sequence[Attachment]: +) -> list[Attachment]: """ Gets all attachments for the current state of the conversation, updating the cache as needed. """ @@ -238,10 +239,14 @@ async def get_attachments( # get all files in the conversation files_response = await context.list_files() + # delete cached attachments that are no longer in the conversation + filenames = {file.filename for file in files_response.files} + asyncio.create_task(_delete_attachments_not_in(context, filenames)) + attachments = [] # for all files, get the attachment for file in files_response.files: - if include_filenames is not None and file.filename not in include_filenames: + if include_filenames and file.filename not in include_filenames: continue if file.filename in exclude_filenames: continue @@ -249,10 +254,6 @@ async def get_attachments( attachment = await _get_attachment_for_file(context, file, {}, error_handler, summarizer=summarizer) attachments.append(attachment) - # delete cached attachments that are no longer in the conversation - filenames = {file.filename for file in files_response.files} - await _delete_attachments_not_in(context, filenames) - return attachments diff --git a/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/archive/__init__.py b/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/archive/__init__.py index faae891d..2b4b3197 100644 --- a/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/archive/__init__.py +++ b/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/archive/__init__.py @@ -2,8 +2,9 @@ Provides the ArchiveTaskQueues class, for integrating with the chat context toolkit's archiving functionality. """ -from ._archive import ArchiveTaskQueues +from ._archive import ArchiveTaskQueues, construct_archive_summarizer __all__ = [ "ArchiveTaskQueues", + "construct_archive_summarizer", ] diff --git a/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/archive/_archive.py b/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/archive/_archive.py index c6952992..36b3f1c4 100644 --- a/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/archive/_archive.py +++ b/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/archive/_archive.py @@ -3,11 +3,12 @@ from chat_context_toolkit.archive import ArchiveReader, ArchiveTaskConfig, ArchiveTaskQueue, StorageProvider from chat_context_toolkit.archive import MessageProvider as ArchiveMessageProvider from chat_context_toolkit.archive.summarization import LLMArchiveSummarizer, LLMArchiveSummarizerConfig -from chat_context_toolkit.history.tool_abbreviations import ToolAbbreviations from openai_client import OpenAIRequestConfig, ServiceConfig, create_client from openai_client.tokens import num_tokens_from_messages from semantic_workbench_assistant.assistant_app import ConversationContext, storage_directory_for_context +from assistant_extensions.attachments._model import Attachment + from ..message_history import chat_context_toolkit_message_provider_for @@ -46,21 +47,30 @@ async def list_files(self, relative_directory_path: PurePath) -> list[PurePath]: def archive_message_provider_for( - context: ConversationContext, service_config: ServiceConfig, request_config: OpenAIRequestConfig + context: ConversationContext, + attachments: list[Attachment], ) -> ArchiveMessageProvider: """Create an archive message provider for the provided context.""" return chat_context_toolkit_message_provider_for( context=context, - tool_abbreviations=ToolAbbreviations(), - service_config=service_config, - request_config=request_config, + attachments=attachments, ) -def _archive_task_queue_for( - context: ConversationContext, +def construct_archive_summarizer( service_config: ServiceConfig, request_config: OpenAIRequestConfig, +) -> LLMArchiveSummarizer: + return LLMArchiveSummarizer( + client_factory=lambda: create_client(service_config), + llm_config=LLMArchiveSummarizerConfig(model=request_config.model), + ) + + +def _archive_task_queue_for( + context: ConversationContext, + attachments: list[Attachment], + archive_summarizer: LLMArchiveSummarizer, archive_task_config: ArchiveTaskConfig = ArchiveTaskConfig(), token_counting_model: str = "gpt-4o", archive_storage_sub_directory: str = "archives", @@ -71,13 +81,11 @@ def _archive_task_queue_for( return ArchiveTaskQueue( storage_provider=ArchiveStorageProvider(context=context, sub_directory=archive_storage_sub_directory), message_provider=archive_message_provider_for( - context=context, service_config=service_config, request_config=request_config + context=context, + attachments=attachments, ), token_counter=lambda messages: num_tokens_from_messages(messages=messages, model=token_counting_model), - summarizer=LLMArchiveSummarizer( - client_factory=lambda: create_client(service_config), - llm_config=LLMArchiveSummarizerConfig(model=request_config.model), - ), + summarizer=archive_summarizer, config=archive_task_config, ) @@ -93,8 +101,8 @@ def __init__(self) -> None: async def enqueue_run( self, context: ConversationContext, - service_config: ServiceConfig, - request_config: OpenAIRequestConfig, + attachments: list[Attachment], + archive_summarizer: LLMArchiveSummarizer, archive_task_config: ArchiveTaskConfig = ArchiveTaskConfig(), ) -> None: """Get the archive task queue for the given context, creating it if it does not exist.""" @@ -102,8 +110,8 @@ async def enqueue_run( if context_id not in self._queues: self._queues[context_id] = _archive_task_queue_for( context=context, - service_config=service_config, - request_config=request_config, + attachments=attachments, + archive_summarizer=archive_summarizer, archive_task_config=archive_task_config, ) await self._queues[context_id].enqueue_run() diff --git a/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/message_history/__init__.py b/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/message_history/__init__.py index 74008453..c85daf66 100644 --- a/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/message_history/__init__.py +++ b/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/message_history/__init__.py @@ -2,8 +2,9 @@ Provides a message history provider for the chat context toolkit's history management. """ -from ._history import chat_context_toolkit_message_provider_for +from ._history import chat_context_toolkit_message_provider_for, construct_attachment_summarizer __all__ = [ "chat_context_toolkit_message_provider_for", + "construct_attachment_summarizer", ] diff --git a/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/message_history/_history.py b/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/message_history/_history.py index 098f4add..8e3e47fb 100644 --- a/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/message_history/_history.py +++ b/libraries/python/assistant-extensions/assistant_extensions/chat_context_toolkit/message_history/_history.py @@ -22,7 +22,7 @@ ) from semantic_workbench_assistant.assistant_app import ConversationContext -from assistant_extensions.attachments._attachments import get_attachments +from assistant_extensions.attachments._model import Attachment from assistant_extensions.attachments._summarizer import LLMConfig, LLMFileSummarizer from ._message import conversation_message_to_chat_message_param @@ -132,11 +132,23 @@ class CompositeMessageProtocol(HistoryMessageProtocol, ArchiveMessageProtocol, P ... -def chat_context_toolkit_message_provider_for( - context: ConversationContext, - tool_abbreviations: ToolAbbreviations, +def construct_attachment_summarizer( service_config: ServiceConfig, request_config: OpenAIRequestConfig, +) -> LLMFileSummarizer: + return LLMFileSummarizer( + llm_config=LLMConfig( + client_factory=lambda: create_client(service_config), + model=request_config.model, + max_response_tokens=request_config.response_tokens, + ) + ) + + +def chat_context_toolkit_message_provider_for( + context: ConversationContext, + attachments: list[Attachment], + tool_abbreviations: ToolAbbreviations = ToolAbbreviations(), ) -> CompositeMessageProvider: """ Create a composite message provider for the given workbench conversation context. @@ -146,9 +158,8 @@ async def provider(after_id: str | None = None) -> Sequence[CompositeMessageProt history = await _get_history_manager_messages( context, tool_abbreviations=tool_abbreviations, - service_config=service_config, - request_config=request_config, after_id=after_id, + attachments=attachments, ) return history @@ -159,8 +170,7 @@ async def provider(after_id: str | None = None) -> Sequence[CompositeMessageProt async def _get_history_manager_messages( context: ConversationContext, tool_abbreviations: ToolAbbreviations, - service_config: ServiceConfig, - request_config: OpenAIRequestConfig, + attachments: list[Attachment], after_id: str | None = None, ) -> list[HistoryMessageWithAbbreviation]: """ @@ -175,21 +185,6 @@ async def _get_history_manager_messages( batch_size = 100 before_message_id = None - attachments = list( - await get_attachments( - context=context, - include_filenames=None, - exclude_filenames=[], - summarizer=LLMFileSummarizer( - llm_config=LLMConfig( - client_factory=lambda: create_client(service_config), - model=request_config.model, - max_response_tokens=request_config.response_tokens, - ) - ), - ) - ) - # each call to get_messages will return a maximum of `batch_size` messages # so we need to loop until all messages are retrieved while True: diff --git a/libraries/python/assistant-extensions/test/attachments/test_attachments.py b/libraries/python/assistant-extensions/test/attachments/test_attachments.py index ed5a2c32..32586f52 100644 --- a/libraries/python/assistant-extensions/test/attachments/test_attachments.py +++ b/libraries/python/assistant-extensions/test/attachments/test_attachments.py @@ -21,6 +21,13 @@ from semantic_workbench_assistant.assistant_app import AssistantAppProtocol, AssistantContext, ConversationContext +@pytest.fixture(scope="function", autouse=True) +def temporary_storage_directory(monkeypatch: pytest.MonkeyPatch) -> Iterable[pathlib.Path]: + with TemporaryDirectory() as tempdir: + monkeypatch.setattr(settings.storage, "root", tempdir) + yield pathlib.Path(tempdir) + + @pytest.mark.parametrize( ("filenames_with_bytes", "expected_messages"), [ @@ -101,7 +108,6 @@ async def test_get_completion_messages_for_attachments( filenames_with_bytes: dict[str, Callable[[], bytes]], expected_messages: list[ChatCompletionMessageParam], - temporary_storage_directory: pathlib.Path, ) -> None: mock_assistant_app = mock.MagicMock(spec=AssistantAppProtocol) @@ -174,10 +180,3 @@ async def read_file_side_effect( ) assert actual_messages == expected_messages - - -@pytest.fixture(scope="function") -def temporary_storage_directory(monkeypatch: pytest.MonkeyPatch) -> Iterable[pathlib.Path]: - with TemporaryDirectory() as tempdir: - monkeypatch.setattr(settings.storage, "root", tempdir) - yield pathlib.Path(tempdir) diff --git a/libraries/python/chat-context-toolkit/CLAUDE.md b/libraries/python/chat-context-toolkit/CLAUDE.md index 15abaee5..28be5e66 100644 --- a/libraries/python/chat-context-toolkit/CLAUDE.md +++ b/libraries/python/chat-context-toolkit/CLAUDE.md @@ -1,4 +1,4 @@ -Python tools: +# Python tools - `uv` for dependency management - `uv run` to run python or project scripts