add search command #157

granawkins · 2023-10-14T07:02:23Z

This PR implements a /search command using embeddings, along with some follow-ups from #144. To use it, add the --use-embedding flag.

biobootloader

Very cool! I have a lot of thoughts, some of these would be follow up PRs:

The first time I ran it I got Embedding batch 1/13... to Embedding batch 13/13... printed out. The second time it just printed Embedding batch 1/1.... I understand the first call had to embed everything, but why did the second call still have a batch? Maybe the batch algorithm always returns at least one batch that might be empty?
I don't think we are tracking embedding costs with the cost logger. I know they are very low, but it might be good to show what they are, so users know? Relatedly, I worry someone will run this on a massive codebase and it'll actually cost them a lot. Like if they have tons of json / data files checked in to the repo or something. Should we have some warning if the number of batches is actually crazy, and it'll cost more than like $1?
What do we do if a file is too big too embed all at once?
When we split embedding sections from whole files to smaller parts it'll be great to show the code in the search results. Formatting that might be tricky. In the VS Code extension we'd make it easy to jump to those files / sections

biobootloader · 2023-10-15T17:39:32Z

mentat/terminal/client.py

@@ -229,7 +229,7 @@ def run_cli():
        help="Exclude the file structure/syntax map from the system prompt",
    )
    parser.add_argument(
-        "--embedding",
+        "--use-embedding",


does --use-embeddings make more sense?

granawkins · 2023-10-19T04:14:51Z

We talked about these but for the record:

The first time I ran it I got Embedding batch 1/13... to Embedding batch 13/13... printed out. The second time it just printed Embedding batch 1/1.... I understand the first call had to embed everything, but why did the second call still have a batch? Maybe the batch algorithm always returns at least one batch that might be empty?

It's embedding the prompt. So if it's a new prompt (hash not in db), there'll always be at least a batch of 1.

I don't think we are tracking embedding costs with the cost logger. I know they are very low, but it might be good to show what they are, so users know? Relatedly, I worry someone will run this on a massive codebase and it'll actually cost them a lot. Like if they have tons of json / data files checked in to the repo or something. Should we have some warning if the number of batches is actually crazy, and it'll cost more than like $1?

I've set this up as you said: give option to ignore embeddings if the cost > $1, otherwise display the cost (with 4 decimal places) afterwards.

What do we do if a file is too big too embed all at once?

These are ignored for now, but I'll make sure they're included when we implement file-splitting.

When we split embedding sections from whole files to smaller parts it'll be great to show the code in the search results. Formatting that might be tricky. In the VS Code extension we'd make it easy to jump to those files / sections

Agreed!

jakethekoenig

Looks good to me just a few small things.

One other note: I notice with small searches, like when I searched for just "parser" the init files score very highly. I wonder if we should not embed very small files. Say ones under 10 characters?

Other than that the searches I tried seemed to surface pretty relevant files.

jakethekoenig · 2023-10-19T18:43:56Z

mentat/commands.py

+            await stream.send(str(e), color="red")
+            return
+
+        for i, (feature, score) in enumerate(results):


I think we should 1 index instead of 0 index.

jakethekoenig · 2023-10-19T18:48:43Z

mentat/commands.py

+
+        for i, (feature, score) in enumerate(results):
+            _i = f"{i}: " if i < 10 else f"{i}:"
+            await stream.send(f"{_i} {score:.3f} | {feature.path}")


You can use :2 to force i to take up the same vertical space:

await stream.send(f"{i:2} {score:.3f} | {feature.path}")

I think that'd be a bit cleaner.

jakethekoenig · 2023-10-19T18:52:11Z

mentat/commands.py

+                await stream.send("\nShow More results? ")
+                if not await ask_yes_no(default_yes=True):
+                    break
+        await stream.send("Search complete", color="green")


I think from a UX perspective it's better not to send this message. The user will know the search is over. We don't send a similar message when they ask a general question and the model responds.

jakethekoenig · 2023-10-19T18:53:16Z

mentat/conversation.py

@@ -171,6 +171,7 @@ async def get_model_response(self) -> list[FileEdit]:
        conversation_history = "\n".join([m["content"] for m in messages_snapshot])
        tokens = count_tokens(conversation_history, self.model)
        response_buffer = 1000
+        print()


Should be removed

jakethekoenig · 2023-10-19T18:57:40Z

mentat/commands.py

+SEARCH_RESULT_BATCH_SIZE = 10
+
+
+class SearchCommand(Command, command_name="search"):


Could we add a SearchCommandTest?

PCSwingle · 2023-10-19T22:30:25Z

mentat/code_context.py

+    ) -> list[tuple[CodeFile, float]]:
+        """Return the top n features that are most similar to the query."""
+        if not self.settings.use_embeddings:
+            raise UserError(


I would prefer if we just sent a red error rather than crashed

PCSwingle

I agree with all of Jake's comments and left one of my own, but looks good to merge after all of that is done! Thanks for adding this!

granawkins · 2023-10-21T03:58:27Z

Thanks for the great feedback, I think I hit everything.

One other note: I notice with small searches, like when I searched for just "parser" the init files score very highly. I wonder if we should not embed very small files. Say ones under 10 characters?

Hmm.. the git_root relative path is included in the embedding. I think returning empty files, as it did, will be useful in some cases. Moving this convo to slack.

granawkins added 2 commits October 14, 2023 13:49

add search command

cc6d3eb

tweaks

0cd1970

biobootloader reviewed Oct 15, 2023

View reviewed changes

fixes from @biobootloader feedback

915cd09

merge in main

23c39e5

granawkins requested review from biobootloader, jakethekoenig and PCSwingle October 19, 2023 06:24

jakethekoenig approved these changes Oct 19, 2023

View reviewed changes

PCSwingle reviewed Oct 19, 2023

View reviewed changes

granawkins added 4 commits October 21, 2023 08:23

merge in main once more

96da1a0

feedback from @jakethekoenig and @PCSwingle

8071836

get context_benchmark working with session_context

5325f38

remove extra imports

33f3dd5

remove unused mock_cost_tracker

d647197

granawkins merged commit 60b7458 into main Oct 21, 2023
8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

add search command #157

add search command #157

granawkins commented Oct 14, 2023

biobootloader left a comment

biobootloader Oct 15, 2023

granawkins commented Oct 19, 2023

jakethekoenig left a comment •

edited

Loading

jakethekoenig Oct 19, 2023

jakethekoenig Oct 19, 2023

jakethekoenig Oct 19, 2023

jakethekoenig Oct 19, 2023

jakethekoenig Oct 19, 2023

PCSwingle Oct 19, 2023

PCSwingle left a comment

granawkins commented Oct 21, 2023 •

edited

Loading

		SEARCH_RESULT_BATCH_SIZE = 10


		class SearchCommand(Command, command_name="search"):

add search command #157

add search command #157

Conversation

granawkins commented Oct 14, 2023

biobootloader left a comment

Choose a reason for hiding this comment

biobootloader Oct 15, 2023

Choose a reason for hiding this comment

granawkins commented Oct 19, 2023

jakethekoenig left a comment • edited Loading

Choose a reason for hiding this comment

jakethekoenig Oct 19, 2023

Choose a reason for hiding this comment

jakethekoenig Oct 19, 2023

Choose a reason for hiding this comment

jakethekoenig Oct 19, 2023

Choose a reason for hiding this comment

jakethekoenig Oct 19, 2023

Choose a reason for hiding this comment

jakethekoenig Oct 19, 2023

Choose a reason for hiding this comment

PCSwingle Oct 19, 2023

Choose a reason for hiding this comment

PCSwingle left a comment

Choose a reason for hiding this comment

granawkins commented Oct 21, 2023 • edited Loading

jakethekoenig left a comment •

edited

Loading

granawkins commented Oct 21, 2023 •

edited

Loading