Add rag enrichment and ingestion. #33413
Conversation
claudevdm
force-pushed
the
add-rag-ingestion-enrichment
branch
2 times, most recently
from
e50d908 to
e966cb7
Compare
claudevdm
force-pushed
the
add-rag-ingestion-enrichment
branch
6 times, most recently
from
ba5ef2a to
b759659
Compare
claudevdm
force-pushed
the
add-rag-ingestion-enrichment
branch
from
b759659 to
b2ad7ac
Compare
|
R: @damccorm |
claudevdm
force-pushed
the
add-rag-ingestion-enrichment
branch
from
b2ad7ac to
589023f
Compare
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
| embedding: ARRAY<FLOAT64> # Vector embedding | ||
| content: STRING # Document content | ||
| metadata: ARRAY<STRUCT< # Nested repeated metadata |
There was a problem hiding this comment.
| metadata: ARRAY<STRUCT< # Nested repeated metadata | |
| metadata: ARRAY<STRUCT> # Nested repeated metadata |
Nit
| ... # chunk.metadata['language'] | ||
| ... metadata_restriction_template="language = '{language}'" | ||
| ... ) | ||
| >>> # When processing a chunk with metadata={'language': 'en'}, |
There was a problem hiding this comment.
I think this could be described a bit more - specifically, I don't think its clear just looking at this object what a Chunk is, and that it will be used by the Enrichment transform/BigQueryVectorSearchEnrichmentHandler.
Maybe at the top of this class, we could mention those 2 concepts (with pydoc links)?
| ... metadata_restriction_template=( | ||
| ... "check_metadata(metadata, 'language', '{language}')" | ||
| ... ) |
There was a problem hiding this comment.
This should just be a string, right?
| ... metadata_restriction_template=( | |
| ... "check_metadata(metadata, 'language', '{language}')" | |
| ... ) | |
| ... metadata_restriction_template="check_metadata(metadata, 'language', '{language}')" |
| embedding_column: Column name containing the embedding vectors. | ||
| columns: List of columns to retrieve from matched vectors. | ||
| neighbor_count: Number of similar vectors to return (top-k). | ||
| metadata_restriction_template: Template string for filtering vectors. |
There was a problem hiding this comment.
I think composing multiple restrictions should work fine, right? It would basically be: metadata_restriction_template="check_metadata(metadata, 'language', '{language}') AND owner = '{companyName}'".
If yes, it would be good to document it. We could also consider taking a list of restrictions, though it is less expressive since it only supports AND, so maybe not worth doing
| """ | ||
| def __init__( | ||
| self, | ||
| project: str, |
There was a problem hiding this comment.
Nit: If we're putting the table name in vector_search_parameters, from a user's perspective I think it makes sense to include the project id as well (even if we just immediately extract that into its own variable in init). Batching kwargs seem fine to keep out since its more of a beam detail than a BQ one.
| - Setting up database connection/client | ||
| - Writing with appropriate batching/error handling | ||
| """ | ||
| pass |
There was a problem hiding this comment.
Nit: It is probably better to throw here like
- it will provide more intelligible errors when a user messes something up (e.g. a typo in the override)| ... ) | ||
| """ | ||
| @abstractmethod | ||
| def create_write_transform(self) -> beam.PTransform: |
There was a problem hiding this comment.
Can we add types to this PTransform? Even if it is beam.PTransform[Chunk, Any]
There was a problem hiding this comment.
Same comment applies below as well, and maybe elsewhere
| @@ -0,0 +1,304 @@ | |||
| # | |||
There was a problem hiding this comment.
Could we call this out in CHANGES.md (both the enrichment and ingestion pieces)?
Add ingestion and enrichment components for rag pipelines.
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.