refactor: ChromaDB migration #20

Anush008 · 2023-08-08T07:56:26Z

Description

This PR intends to migrate the project to use ChromaDB by implementing the RepositoryEmbeddingsDB trait for the ChromaDB client.

repo-query/src/db/mod.rs

Lines 9 to 22 in 117a933

 pub trait RepositoryEmbeddingsDB { 

 fn insert_repo_embeddings(&self, repo: RepositoryEmbeddings) -> Result<()>; 

 fn get_relevant_file_paths( 

 &self, 

 repository: &Repository, 

 query_embeddings: Embeddings, 

 limit: usize, 

 ) -> Result<RepositoryFilePaths>; 

 fn get_file_paths(&self, repository: &Repository) -> Result<RepositoryFilePaths>; 

 fn is_indexed(&self, repository: &Repository) -> Result<bool>; 

 }

The README has been updated with the new local development steps
The ChromaDB service has been added to the docker-compose.yaml file.

What type of PR is this? (check all applicable)

Related Tickets & Documents

Resolves #19

Mobile & Desktop Screenshots/Recordings

Added tests?

👍 yes
🙅 no, because they aren't needed
🙋 no, because I need help

Added to documentation?

Anush008 · 2023-08-08T13:38:57Z

@bdougie, @jpmcb, you can try the service now.
Everything should stay the same.
You can set the OPENAI_API_KEY and WEBSERVER_PORT values in a .env file and run make up quickly get this running.

jpmcb · 2023-08-08T14:38:39Z

Wow excellent work. I'll review this today and look at getting this deployed on our azure instance. I have some terraform that should make that a quick process.

Anush008 · 2023-08-08T15:08:58Z

Hey @jpmcb. The ort lib supports some execution providers that can improve the embedding performance.
https://github.com/pykeio/ort#execution-providers

We are currently using the CPU provider.

repo-query/src/embeddings/onnx.rs

Lines 22 to 24 in afc4d19

 .with_execution_providers([ExecutionProvider::CPU( 

 CPUExecutionProviderOptions::default(), 

 )])

.

Moving forward, if needed, we can try some of the others if the hardware support is available.

jpmcb · 2023-08-08T17:18:58Z

if the hardware support is available

GPU support would be cool, but right now, for the way we have it deployed through Container Apps, there aren't GPU instances available. This is essentially their serverless container runtime and doesn't involve alot of infrastructure. We'd need to look at a more intensive deployment strategy, maybe through Container Instances or Azure's Kubernetes service, that would let us deploy GPU instances.

For now, until we need to start scaling up in that way, I think the Container Apps deployment with CPU ort optimization is just fine.

bdougie

Looks like a pretty clean swap. Really appreciate you kicking off the chroma client work as well.

jpmcb

I was having some snags running the image locally on arm64 so i opened: chroma-core/chroma#972

But I was able to rebuild the image and this works well - nice work!

Performance seems slightly worse than qdrant locally for me. But will have to try that on real hardware.

jpmcb · 2023-08-12T14:40:28Z

src/db/mod.rs

+ fn insert_repo_embeddings(&self, repo: RepositoryEmbeddings) -> Result<()>;

- async fn get_relevant_files(
+ fn get_relevant_file_paths(
 &self,
 repository: &Repository,
 query_embeddings: Embeddings,
 limit: usize,
 ) -> Result<RepositoryFilePaths>;

- async fn get_file_paths(&self, repository: &Repository) -> Result<RepositoryFilePaths>;
+ fn get_file_paths(&self, repository: &Repository) -> Result<RepositoryFilePaths>;

- async fn is_indexed(&self, repository: &Repository) -> Result<bool>;
+ fn is_indexed(&self, repository: &Repository) -> Result<bool>;


Makes sense we'd remove async since before we were just await on each call anyways 👍🏼

jpmcb · 2023-08-12T15:33:16Z

.env.example

+OPENAI_API_KEY= #REQUIRED
+WEBSERVER_PORT= #REQUIRED
+CHROMA_URL= #Defaults to http://localhost:8000


Where does this get used to set the Chroma client? I'm not seeing any other references to this in the code. Is it being consumed by the Chroma rust client?

Yes. The Chromadb-rs lib uses the CHROMA_URL env by default.

repo-query/src/db/chroma.rs

Line 21 in 37bd468

let client = ChromaClient::new(Default::default());

Anush008 · 2023-08-13T09:54:04Z

cc7f606, resolves #21.

jpmcb

There are a few non-starters to getting this merged:

[Feature Request]: Remove re-build during docker container entry chroma-core/chroma#1002
- Rebuilding during container entry causes quite a lengthy gap for the repo-query service to be ready. Ideally, since we're deploying all this in a conainer app environment, we'd not have to wait for container bits to be rebuilt
- [Feature Request]: Provide a precompiled version of chroma-hnswlib chroma-core/chroma#993
  - If we had a precompiled form of the chroma database that we could drop into our own container, that would alleviate this pain, so we may want to watch that issue as well
[Bug]: ghcr.io/chroma-core/chroma does not support arm64 architecture chroma-core/chroma#972
- The team is working on multi-arch builds and that will help alot with local development for our contributor community

Until then, we should hold this migration and stick with Qdrant since it's working out of the box in our container environments.

Anush008 added 5 commits August 8, 2023 11:16

refactor: migrate to ChromaDB

3c705d0

build: add chromadb service

6d1a614

docs: Update README.md

117a933

chore: CHROMA_URL service name

343c2bf

chore: Removed default webserver port fallback

e258260

Anush008 marked this pull request as ready for review August 8, 2023 13:38

Anush008 requested review from jpmcb and bdougie August 8, 2023 13:47

bdougie approved these changes Aug 9, 2023

View reviewed changes

chore: Update deps

37bd468

jpmcb reviewed Aug 12, 2023

View reviewed changes

docs: Update API reference README.md

cc7f606

Anush008 linked an issue Aug 13, 2023 that may be closed by this pull request

Collection example in the README seems off #21

Closed

jpmcb requested changes Aug 17, 2023

View reviewed changes

Anush008 closed this Sep 7, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

refactor: ChromaDB migration #20

refactor: ChromaDB migration #20

Anush008 commented Aug 8, 2023 •

edited

Anush008 commented Aug 8, 2023 •

edited

jpmcb commented Aug 8, 2023

Anush008 commented Aug 8, 2023 •

edited

jpmcb commented Aug 8, 2023

bdougie left a comment

jpmcb left a comment •

edited

jpmcb Aug 12, 2023

jpmcb Aug 12, 2023

Anush008 Aug 12, 2023 •

edited

Anush008 commented Aug 13, 2023

jpmcb left a comment

	pub trait RepositoryEmbeddingsDB {
	fn insert_repo_embeddings(&self, repo: RepositoryEmbeddings) -> Result<()>;

	fn get_relevant_file_paths(
	&self,
	repository: &Repository,
	query_embeddings: Embeddings,
	limit: usize,
	) -> Result<RepositoryFilePaths>;

	fn get_file_paths(&self, repository: &Repository) -> Result<RepositoryFilePaths>;

	fn is_indexed(&self, repository: &Repository) -> Result<bool>;
	}

refactor: ChromaDB migration #20

refactor: ChromaDB migration #20

Conversation

Anush008 commented Aug 8, 2023 • edited

Description

What type of PR is this? (check all applicable)

Related Tickets & Documents

Mobile & Desktop Screenshots/Recordings

Added tests?

Added to documentation?

Anush008 commented Aug 8, 2023 • edited

jpmcb commented Aug 8, 2023

Anush008 commented Aug 8, 2023 • edited

jpmcb commented Aug 8, 2023

bdougie left a comment

Choose a reason for hiding this comment

jpmcb left a comment • edited

Choose a reason for hiding this comment

jpmcb Aug 12, 2023

Choose a reason for hiding this comment

jpmcb Aug 12, 2023

Choose a reason for hiding this comment

Anush008 Aug 12, 2023 • edited

Choose a reason for hiding this comment

Anush008 commented Aug 13, 2023

jpmcb left a comment

Choose a reason for hiding this comment

Anush008 commented Aug 8, 2023 •

edited

Anush008 commented Aug 8, 2023 •

edited

Anush008 commented Aug 8, 2023 •

edited

jpmcb left a comment •

edited

Anush008 Aug 12, 2023 •

edited