Skip to content

Add Chroma Retrieval #323 #324

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add Chroma Retrieval #323 #324

wants to merge 1 commit into from

Conversation

ailunc
Copy link

@ailunc ailunc commented May 7, 2025

Issue Link (REQUIRED)

#323
Fixes #323

Summary

Changes

  • Added Python implementation of Chroma Retriever
    • Created ChromaRetriever class and ChromaRetrieverOptions class
    • Implemented all necessary retrieval methods
    • Added unit tests
    • Updated dependencies and exports
  • Added TypeScript implementation of Chroma Retriever
    • Created ChromaRetriever class and ChromaRetrieverOptions interface
    • Implemented all necessary retrieval methods
    • Added unit tests
    • Updated dependencies and exports
    • Configured testing environment

User experience

Before:

  • Users could not use ChromaDB as a retriever
  • Manual implementation of ChromaDB integration was required

After:

  • Users can directly use Chroma Retriever
  • Support for both local and remote ChromaDB instances
  • Support for custom embedding functions
  • Support for persistent storage
  • Support for SSL connections
  • Support for custom result count
  • Support for metadata retrieval

Usage examples:

Python:

from agent_squad.retrievers import ChromaRetriever, ChromaRetrieverOptions

# Create retriever
options = ChromaRetrieverOptions(
    collection_name="my_collection",
    persist_directory="./data"
)
retriever = ChromaRetriever(options)

# Retrieve documents
results = await retriever.retrieve("query text")
combined_results = await retriever.retrieve_and_combine_results("query text")

TypeScript:

import { ChromaRetriever, ChromaRetrieverOptions } from 'agent-squad';

// Create retriever
const options: ChromaRetrieverOptions = {
    collectionName: "my_collection",
    persistDirectory: "./data"
};
const retriever = new ChromaRetriever(options);

// Retrieve documents
const results = await retriever.retrieve("query text");
const combinedResults = await retriever.retrieveAndCombineResults("query text");

Checklist

  • I have performed a self-review of this change
  • Changes have been tested
  • Changes are documented
  • I have linked this PR to an existing issue (required)
Is this a breaking change?

RFC issue number: N/A

Checklist:

  • Migration process documented
  • Implement warnings (if it can live side by side)

Acknowledgment

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Disclaimer: We value your time and bandwidth. As such, any pull requests created on non-triaged issues might not be successful.

Copy link

@jeffchuber jeffchuber left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks for this!!! added a very clarification requests/comments

port: Optional[int] = None
ssl: Optional[bool] = False
n_results: Optional[int] = 4
embedding_function: Optional[Any] = None

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should this be typed? also im thinking it shouldnt be optional?

host: Optional[str] = None
port: Optional[int] = None
ssl: Optional[bool] = False
n_results: Optional[int] = 4

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggest higher default value like 10?

results = await self.retrieve(text)
return self.combine_retrieval_results(results)

async def retrieve_and_generate(self, text: str) -> Any:

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what does this do and is there a reason it cant be implemented?

@@ -40,6 +40,7 @@
"@aws-sdk/util-dynamodb": "^3.621.0",
"@libsql/client": "^0.14.0",
"axios": "^1.7.2",
"chromadb": "^1.8.1",

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why did you set 1.8.1 here? the latest version is 2.4.3

@@ -49,6 +50,7 @@
"devDependencies": {
"@types/jest": "^29.5.12",
"@types/mocha": "^10.0.7",
"@types/node": "^20.11.30",

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

was this a required change?

port?: number;
ssl?: boolean;
nResults?: number;
embeddingFunction?: any;

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same as python - should this be required and typed?


const results = await this.collection.query({
queryTexts: [text],
nResults: this.options.nResults || 4

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

bump default results to 10?

}

/**
* Placeholder for retrieve and generate functionality

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same question as python

setup(
install_requires=[

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

will the maintainers be ok with this?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Feature request: Add Chroma Retrieval
2 participants