[ENH] JS Client Refactor Part 1 - Refactoring the Client

[AlabasterAxe]

Description of changes

The changes in this PR are described in detail in the "Phase 1" section of this doc: https://docs.google.com/document/d/1DH9Gl5UQeuDr6vyGgoYUw0TkBxITgXIutZSU-ZYdP_E/edit#heading=h.wnaf3x8gqzzl

The TL;DR is that this PR consists of the following changes:

• turns Collection into a data object (with the exception of the embedding function) and flattens its methods into the top-level ChromaClient
• Improves the typing of the methods on the client so that we do a better job of handling single vs multiple cases.
• validates tenant and database on client startup
• removes some unnecessary code in the client that was just used to enable certain test assertions

This PR corresponds with this issue: #2333

This is the first PR in a series that comprise the whole client refactor:

• Part 2: Examples
• Part 3: Tests
• Part 4: Docs
• Full Stack PR

Test plan

• all unit tests are passing

Documentation Changes

• will need to update docs as this is a breaking change with the previous API

[github-actions]

Reviewer Checklist

Please leverage this checklist to ensure your code review is thorough before approving

Testing, Bugs, Errors, Logs, Documentation

• Can you think of any use case in which the code does not behave as intended? Have they been tested?
• Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
• If appropriate, are there adequate property based tests?
• If appropriate, are there adequate unit tests?
• Should any logging, debugging, tracing information be added or removed?
• Are error messages user-friendly?
• Have all documentation changes needed been made?
• Have all non-obvious changes been commented?

System Compatibility

• Are there any potential impacts on other parts of the system or backward compatibility?
• Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?

Quality

• Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)

[AlabasterAxe]

Instead of having queryTexts and queryEmbeddings separately, I'm thinking we can just have this single type DocQuery which can be a string, which will be converted to an embedding, according to the configured embedding function in the collection.

[codetheweb]

not necessarily opposed to this, but wondering how much we want to try to keep Python/JS interfaces roughly in sync

@atroyn wdyt

[atroyn]

I think for this PR, we want to keep them in sync. This might be a bit janky for now but I would prefer that we review the interface as a whole between both clients rather than changing it piecemeal in each.

Can we create a new issue to discuss this?

[AlabasterAxe]

#2483

[AlabasterAxe]

@atroyn does that apply to the pluralization logic as well? I've added convenience overrides for singular document operations e.g:

const resp = getRecords(collection, {
  id: "blah"
});

console.log(resp); // prints "{id: "blah", embedding: [1, 2, 3], document: "whatever" }"

[AlabasterAxe]

The way I've approached this in similar situations is to have a shared init promise that all methods await at the beginning. Since all of the methods on this class are async anyway, we just asynchronously validate that it's been initted successfully when the user calls any of the methods.

[codetheweb]

interesting, wondering if it's worth wrapping the class with a Proxy so we don't have to manually add it to each public method?

[AlabasterAxe]

Yeah that is a good callout. I think there's a high likelihood that this will get dropped in the future. One thing that the current approach allows is to put certain validations before the await initPromise; e.g.

if (typeof params.ids === "string") {
  throw new Error(
    "To add a single record, use the singular field names (id, embedding, metadata, document) instead of the plural field names (ids, embeddings, metadatas, documents)",
  );
}
await this.initPromise;

Also, ideally, this would only apply to async functions. Do you know if there is a way to detect that in a proxy and only wrap the async functions? Obviously we only have async functions today but I would worry about accidentally making a synchronous function async.

[AlabasterAxe]

But yeah not opposed to this. My gut reaction would be to hold off on this, but happy to do this in this PR if you think that's the better trade-off.

[codetheweb]

I think we can leave as-is for now, easy to adjust later

[atroyn]

Adding @codetheweb for review; will cover this this week!

[codetheweb]

not necessarily opposed to this, but wondering how much we want to try to keep Python/JS interfaces roughly in sync

@atroyn wdyt

[codetheweb]

interesting, wondering if it's worth wrapping the class with a Proxy so we don't have to manually add it to each public method?

[codetheweb]

most (all?) methods that deal with documents/records only need the collection's ID, so maybe we could update the typing to reflect that?

Suggested change

	collection: Collection,
	collection: Pick<Collection, "id">,

then people can avoid an API call to fetch full collection details if they really want to

[AlabasterAxe]

Yeah the one thing that we do regularly need is the embedding function so I think it'll beneficial for us to keep Collection as a kind of opaque object that they have to get from the API. My understanding is that there's work underway for the backend to store some representation of embedding function alongside the collection which I think will make this opaque Collection object make more sense.

@atroyn in case I've misunderstood anything here.

[codetheweb]

personally, as a JS dev, I would assume any object returned from an idiomatic API client library is trivially serializable/de-serializable (JSON.[stringify|parse])
as long as we can hold that assumption that sounds ok to me

[codetheweb]

alternatively, if every single method only needs the collection ID, I do feel it would be more idiomatic to have methods accept async addRecords(collection_id: string...) instead (and store any necessary collection metadata on the client object)

[atroyn]

We will be able to hold that assumption with EF persist in the works

[codetheweb]

@AlabasterAxe what do you think about using the collection ID for these methods?

[AlabasterAxe]

Okay so if I'm understanding what you proposed above, you're saying that the client would maintain a mapping from collection ids to embedding functions and then will just look that up whenever the user supplies a collection id?

One potential issue that could arise is that a user could attempt to call the addRecords method even if the embedding function had never been configured for a particular collection.

So hypothetically we could have the following situation:

/* Run 1 */
const collection = client.createCollection({name: "test"});
console.log(collection.id) // prints 1234

/* Run 2 */
client.addRecords("1234", {/* record contents */}); // How do we know which embedding function to use 
                                                    // here if the client has not gotten the collection in this
                                                    // run of the script?

I'll defer to @atroyn here: Should we continue accepting collection objects in the record methods? or should we move to accepting collection ids in the API?

[codetheweb]

discussed with @atroyn and we settled on keeping the methods accepting collection objects for now

[codetheweb]

personally, as a JS dev, I would assume any object returned from an idiomatic API client library is trivially serializable/de-serializable (JSON.[stringify|parse])
as long as we can hold that assumption that sounds ok to me

[atroyn]

Initial review pass looks good, @codetheweb is reviewing docs currently.

@AlabasterAxe could we please split this up into a stack of PRs for easier in-depth review of each part?

[codetheweb]

docs look pretty good to me 👍

[AlabasterAxe]

Initial review pass looks good, @codetheweb is reviewing docs currently.

@AlabasterAxe could we please split this up into a stack of PRs for easier in-depth review of each part?

@atroyn I don't have access to Chroma's Graphite instance so I did this manually. Let me know if that works for you. Also this PR is in kind of a weird state. I don't have permission to push branches to the chroma repo directly so I started it from my personal chroma fork. That means the subsequent PRs are not considered as being against the chroma repo yet but I think that should be okay because once you all give the go ahead I can merge them in quick succession.

Also, it sounds like it may be good to sync about all of the behavior changes this PR includes because it sounds like I may have diverged a bit too much from the Python client. Let me know if you'd like to sync quickly and then I can make the final changes to this PR so it's ready for the final round of reviews.

[AlabasterAxe]

@atroyn

I've recorded my thoughts about the various suggested changes to client behavior in these isues:
#2483 (the queryTexts vs query issue from above)
#2492
#2493
#2494

As discussed yesterday these are just for posterity/a jumping off point for larger conversations about client behavior as opposed to proposals that I'm advocating we implement in the near term. I still need to revert the behavior changes as discussed from yesterday though. Will ping on this PR when those changes are in.

[AlabasterAxe]

@atroyn Okay, I've made the requested changes so these PRs should be ready for another look. Specifically:

• revert the merging of queryTexts and queryEmbedding into a single query param
• revert the singular field naming
• revert the parameter-dependent return type (i.e. if user supplies a singular type, unwrap resulting array)

Also, I introduced a some style regressions along the way so I've just added a final formatting PR at the end of the stack.

[codetheweb]

@AlabasterAxe what do you think about using the collection ID for these methods?

[codetheweb]

validates tenant and database on client startup

let's extract this out into a separate PR, we can merge it before this entire stack is approved since it doesn't change any interfaces

also, a quick PR to export types in package.jsoncould be merged before this stack is approved as well

otherwise this first PR is looking good! (outside of failing checks)

[codetheweb]

does this signature (and others like it below) need to be overloaded since the response shape doesn't change?

[AlabasterAxe]

@codetheweb validation PR here: #2537

This PR is rebased on top of that PR but I'm still using the chroma repo main as its base for this PR.

[AlabasterAxe]

@codetheweb

I've propagated all of the PR feedback through the stack and put together the "Full stack" PR here: #2542

All PR checks are passing there so I'm thinking that as long as you approve of each of the PRs in the stack, we should be good to merge the Full Stack PR (even though e.g. this PR is failing checks) WDYT?

[AlabasterAxe]

NOTE: This check has been removed in a subsequent stack (here and below). Not changing here to avoid further PR churn.