Conversation get_messages gets code_message #533

jakethekoenig · 2024-02-23T19:31:29Z

Frequently in the codebase people will go to conversation for the messages and then independently get the code_message. It's simpler if the conversation handles that.

get_messages is also changed to accept a system prompt argument instead of having a include_system_prompt parameter to simplify agent's calling of it. If callers genuinely want no system prompt they can pass [].

Move responsibility to raise if context too big from code_context to conversation.

Pull Request Checklist

Documentation has been updated, or this change doesn't require that

Frequently in the codebase people will go to conversation for the messages and then independently get the code_message. It's simpler if the conversation handles that. The sampler and reviser both do this but with a different system prompt. I think we should also change get_conversation to have an optional system prompt parameter that overrides the default parser system prompt when passed. We should also provide a get_token_count method as many callers are getting the messages just for this one piece of information.

mentatbot · 2024-02-23T19:40:39Z

mentat/agent_handler.py

@@ -92,7 +92,7 @@ async def _determine_commands(self) -> List[str]:
            ChatCompletionSystemMessageParam(


Great job making the get_messages call asynchronous. This ensures consistency with the rest of the codebase.

mentatbot · 2024-02-23T19:40:41Z

mentat/code_context.py

@@ -68,7 +66,7 @@ def __init__(
        self.ignore_files: Set[Path] = set()
        self.auto_features: List[CodeFeature] = []



Making refresh_context_display asynchronous is a good move for consistency and potentially allows for more complex operations within the method in the future.

mentatbot · 2024-02-23T19:40:43Z

mentat/conversation.py

@@ -45,20 +45,7 @@ async def display_token_count(self):
        config = session_context.config


The addition of include_code_message parameter in get_messages method is a smart way to reduce redundancy. It's a good practice to encapsulate this logic within the conversation class.

mentatbot · 2024-02-23T19:40:44Z

mentat/conversation.py

    ) -> list[ChatCompletionMessageParam]:
        """Returns the messages in the conversation. The system message may change throughout
        the conversation and messages may contain additional metadata not supported by the API,
        so it is important to access the messages through this method.
        """
-        session_context = SESSION_CONTEXT.get()
-        config = session_context.config
+        ctx = SESSION_CONTEXT.get()


The async keyword addition here is crucial for the method's functionality with the new async calls inside it. Well done on updating the method signature accordingly.

mentatbot · 2024-02-23T19:40:46Z

mentat/revisor/revisor.py

@@ -54,7 +54,7 @@ async def revise_edit(file_edit: FileEdit):
    user_message = list(


Updating to the async get_messages call here ensures that the revisor's behavior remains consistent with the rest of the system's async nature.

mentatbot · 2024-02-23T19:40:48Z

mentat/sampler/sampler.py

@@ -118,7 +118,8 @@ async def create_sample(self) -> Sample:
        message_history: list[dict[str, str]] = []


It's good to see the async get_messages being used here. This change ensures that the sampler's operations are in line with the async architecture of the application.

mentatbot · 2024-02-23T19:40:50Z

mentat/session.py

@@ -151,7 +151,7 @@ async def _main(self):
        stream.send("Type 'q' or use Ctrl-C to quit at any time.")


The change to call refresh_context_display asynchronously is necessary and correct. It's important for maintaining the non-blocking nature of the application.

chose the wrong line to comment on here?

mentatbot · 2024-02-23T19:40:52Z

mentat/session_input.py

@@ -60,7 +60,7 @@ async def collect_input_with_commands() -> StreamMessage:
            arguments = shlex.split(" ".join(response.data.split(" ")[1:]))


The update to call refresh_context_display asynchronously here is a good catch. It ensures that the command handling process doesn't block other operations.

mentatbot · 2024-02-23T19:40:54Z

tests/commands_test.py

@@ -246,7 +246,8 @@ async def test_clear_command(temp_testbed, mock_collect_user_input, mock_call_ll
    await session.stream.recv(channel="client_exit")


Updating the test to await get_messages is necessary due to the changes in the Conversation class. Good job on keeping the tests up to date.

mentatbot · 2024-02-23T19:40:56Z

tests/conversation_test.py

@@ -1,44 +1,50 @@
+import pytest


The addition of async and await in the tests is essential for testing the new asynchronous get_messages method. It's great to see the tests being adapted to match the code changes.

mentatbot · 2024-02-23T19:40:58Z

MENTAT CODE REVIEW IN BETA. Please Reply with feedback

This pull request introduces several key improvements, notably making various methods asynchronous to align with the overall async architecture of the application. The changes are well thought out and implemented, ensuring consistency and efficiency in message handling and command execution. It's particularly commendable how the changes have been propagated throughout the codebase, including updating tests to reflect the new asynchronous behavior. Overall, this PR represents a solid step forward in maintaining and improving the codebase's quality and functionality.

jakethekoenig · 2024-02-26T17:59:19Z

mentat/code_context.py


        # Calculate user included features token size
        include_features = [
            feature
            for file_features in self.include_files.values()
            for feature in file_features
        ]
-        include_files_message = get_code_message_from_features(include_features)


I moved all this to the if because it's only relevant if auto_context is on. Benchmarking shows its basically unimportant for performance but I think its easier to see how things are used and the flow of logic this way.

granawkins

I like removing all the token checks from the context-building process (moving to llm_api_handler instead), and the conversation.count_tokens method. Made one small suggestion, otherwise LGTM. 🚀

granawkins · 2024-02-26T22:51:11Z

mentat/conversation.py

+        messages_snapshot = await self.get_messages(include_code_message=True)
+        tokens_used = prompt_tokens(messages_snapshot, config.model)


Suggested change

messages_snapshot = await self.get_messages(include_code_message=True)

tokens_used = prompt_tokens(messages_snapshot, config.model)

tokens_used = await self.count_tokens(include_code_message=True)

Can this new method be reused here?

It could be used but we still need to compute messages_snapshot to send to the llm anyway. So I think it is better as is.

mentatbot bot reviewed Feb 23, 2024

View reviewed changes

jakethekoenig added 3 commits February 26, 2024 09:20

System prompt argument to get_messages

6eb7d4a

Add count_tokens to conversation

963b226

Move raise responsibility

5599ca4

jakethekoenig commented Feb 26, 2024

View reviewed changes

granawkins approved these changes Feb 26, 2024

View reviewed changes

jakethekoenig merged commit 7b04721 into main Feb 27, 2024
16 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Conversation get_messages gets code_message #533

Conversation get_messages gets code_message #533

jakethekoenig commented Feb 23, 2024 •

edited

Loading

mentatbot bot Feb 23, 2024

mentatbot bot Feb 23, 2024

mentatbot bot Feb 23, 2024

mentatbot bot Feb 23, 2024

mentatbot bot Feb 23, 2024

mentatbot bot Feb 23, 2024

mentatbot bot Feb 23, 2024

biobootloader Feb 27, 2024

mentatbot bot Feb 23, 2024

mentatbot bot Feb 23, 2024

mentatbot bot Feb 23, 2024

mentatbot bot commented Feb 23, 2024

jakethekoenig Feb 26, 2024

granawkins left a comment

granawkins Feb 26, 2024

jakethekoenig Feb 27, 2024

		@@ -92,7 +92,7 @@ async def _determine_commands(self) -> List[str]:
		ChatCompletionSystemMessageParam(

		@@ -68,7 +66,7 @@ def __init__(
		self.ignore_files: Set[Path] = set()
		self.auto_features: List[CodeFeature] = []

		@@ -45,20 +45,7 @@ async def display_token_count(self):
		config = session_context.config

		@@ -54,7 +54,7 @@ async def revise_edit(file_edit: FileEdit):
		user_message = list(

		@@ -118,7 +118,8 @@ async def create_sample(self) -> Sample:
		message_history: list[dict[str, str]] = []

		@@ -151,7 +151,7 @@ async def _main(self):
		stream.send("Type 'q' or use Ctrl-C to quit at any time.")

		@@ -60,7 +60,7 @@ async def collect_input_with_commands() -> StreamMessage:
		arguments = shlex.split(" ".join(response.data.split(" ")[1:]))

		@@ -246,7 +246,8 @@ async def test_clear_command(temp_testbed, mock_collect_user_input, mock_call_ll
		await session.stream.recv(channel="client_exit")

		messages_snapshot = await self.get_messages(include_code_message=True)
		tokens_used = prompt_tokens(messages_snapshot, config.model)

	messages_snapshot = await self.get_messages(include_code_message=True)
	tokens_used = prompt_tokens(messages_snapshot, config.model)
	tokens_used = await self.count_tokens(include_code_message=True)

Conversation get_messages gets code_message #533

Conversation get_messages gets code_message #533

Conversation

jakethekoenig commented Feb 23, 2024 • edited Loading

Pull Request Checklist

mentatbot bot Feb 23, 2024

Choose a reason for hiding this comment

mentatbot bot Feb 23, 2024

Choose a reason for hiding this comment

mentatbot bot Feb 23, 2024

Choose a reason for hiding this comment

mentatbot bot Feb 23, 2024

Choose a reason for hiding this comment

mentatbot bot Feb 23, 2024

Choose a reason for hiding this comment

mentatbot bot Feb 23, 2024

Choose a reason for hiding this comment

mentatbot bot Feb 23, 2024

Choose a reason for hiding this comment

biobootloader Feb 27, 2024

Choose a reason for hiding this comment

mentatbot bot Feb 23, 2024

Choose a reason for hiding this comment

mentatbot bot Feb 23, 2024

Choose a reason for hiding this comment

mentatbot bot Feb 23, 2024

Choose a reason for hiding this comment

mentatbot bot commented Feb 23, 2024

jakethekoenig Feb 26, 2024

Choose a reason for hiding this comment

granawkins left a comment

Choose a reason for hiding this comment

granawkins Feb 26, 2024

Choose a reason for hiding this comment

jakethekoenig Feb 27, 2024

Choose a reason for hiding this comment

jakethekoenig commented Feb 23, 2024 •

edited

Loading