Skip to content

Enable embedding caching on all vectorizers #320

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

Merged
merged 4 commits into from
Apr 17, 2025

Conversation

tylerhutcherson
Copy link
Collaborator

Adds support to the BaseVectorizer class to have an optional EmbeddingsCache attached.

  • Refactored the subclass vectorizers to implement private embed methods and then let the base class handle the cache wrapper logic.
  • Fixed some circular imports.
  • Fixed async client handling in the cache subclasses (caught during testing).
  • Handle some typing checks and pydantic stuff related to private attrs and custom attrs.

TODO in a separate PR:

  • Add embeddings caching to our testing suite (CI/CD speed up??)
  • Add embeddings caching to our SemanticRouter

@tylerhutcherson tylerhutcherson added the enhancement New feature or request label Apr 16, 2025
@tylerhutcherson tylerhutcherson marked this pull request as ready for review April 16, 2025 22:34
@tylerhutcherson tylerhutcherson changed the title Feat/raae 594 cached vectorizers Add caching to vectorizers Apr 16, 2025
Copy link
Collaborator

@abrookins abrookins left a comment

Choose a reason for hiding this comment

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

This looks great! Nice culmination of a lot of work. Had one non-blocking suggestion to consider, totally optional. 👍


try:
# Efficient batch cache lookup
cache_results = await self.cache.amget(texts=texts, model_name=self.model)
Copy link
Collaborator

Choose a reason for hiding this comment

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

The whole amget / efficiency aspect of this work is nice. 🔥


model: str
dtype: str = "float32"
dims: Optional[int] = None
dims: Annotated[Optional[int], Field(strict=True, gt=0)] = None
cache: Optional[EmbeddingsCache] = Field(default=None)
Copy link
Collaborator

Choose a reason for hiding this comment

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

Passing a cache object in to get caching feels elegant

async def _aembed(self, text: str, **kwargs) -> List[float]:
"""Asynchronously generate a vector embedding for a single text.

Note: This implementation falls back to the synchronous version as
Copy link
Collaborator

Choose a reason for hiding this comment

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

This could be worth logging when these methods are called.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

done!

@tylerhutcherson tylerhutcherson changed the title Add caching to vectorizers Enable embedding caching on all vectorizers Apr 17, 2025
@tylerhutcherson tylerhutcherson merged commit 95ffe75 into 0.6.0 Apr 17, 2025
31 checks passed
@tylerhutcherson tylerhutcherson deleted the feat/RAAE-594-cached-vectorizers branch April 17, 2025 14:21
abrookins pushed a commit that referenced this pull request May 1, 2025
Adds support to the `BaseVectorizer` class to have an optional
`EmbeddingsCache` attached.

- Refactored the subclass vectorizers to implement private embed methods
and then let the base class handle the cache wrapper logic.
- Fixed some circular imports.
- Fixed async client handling in the cache subclasses (caught during
testing).
- Handle some typing checks and pydantic stuff related to private attrs
and custom attrs.

TODO in a separate PR:
- Add embeddings caching to our testing suite (CI/CD speed up??)
- Add embeddings caching to our SemanticRouter
abrookins pushed a commit that referenced this pull request May 1, 2025
Adds support to the `BaseVectorizer` class to have an optional
`EmbeddingsCache` attached.

- Refactored the subclass vectorizers to implement private embed methods
and then let the base class handle the cache wrapper logic.
- Fixed some circular imports.
- Fixed async client handling in the cache subclasses (caught during
testing).
- Handle some typing checks and pydantic stuff related to private attrs
and custom attrs.

TODO in a separate PR:
- Add embeddings caching to our testing suite (CI/CD speed up??)
- Add embeddings caching to our SemanticRouter
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants