You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "ferruzzi (via GitHub)" <gi...@apache.org> on 2024/05/08 21:01:44 UTC
[PR] Amaazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
ferruzzi opened a new pull request, #39500:
URL: https://github.com/apache/airflow/pull/39500
Both of these calls are super fast and neither has any kind of waiter or means of checking the status, so here can not be any sensor or trigger for them. They are simple client calls, but I think making these Operators allowed us to simplify the complicated formatting on the client API call itself, for a better UX.
One question I have is the parameter name for `input` and `query`. I feel like these should be standardized, but also we usually go out of our way to stick as close as possible tot he API call's parameter names. Let me know what you think. For InvokeModel I used "input_data", I could standardize all thee on that name or just leave them as i have here.
All related unit tets, CI, and build-docs pass locally.
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [newsfragments](https://github.com/apache/airflow/tree/main/newsfragments).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1600319579
##########
tests/system/providers/amazon/aws/example_bedrock_knowledge_base.py:
##########
@@ -480,6 +483,24 @@ def delete_opensearch_policies(collection_name: str):
)
# [END howto_sensor_bedrock_ingest_data]
+ # [START howto_operator_bedrock_retrieve_and_generate]
+ retrieve_and_generate = BedrockRaGOperator(
+ task_id="retrieve_and_generate",
+ input="Who was the CEO of Amazon on 2022?",
+ source_type="KNOWLEDGE_BASE",
Review Comment:
Sorry, kinda repeated myself there. The only new news in that last comment was the bit about aiobotocore being a limiting factor, and that I manually tested it in breeze by forcing the upgrade in the breeze env.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas merged PR #39500:
URL: https://github.com/apache/airflow/pull/39500
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1599301829
##########
tests/system/providers/amazon/aws/example_bedrock_knowledge_base.py:
##########
@@ -480,6 +483,24 @@ def delete_opensearch_policies(collection_name: str):
)
# [END howto_sensor_bedrock_ingest_data]
+ # [START howto_operator_bedrock_retrieve_and_generate]
+ retrieve_and_generate = BedrockRaGOperator(
+ task_id="retrieve_and_generate",
+ input="Who was the CEO of Amazon on 2022?",
+ source_type="KNOWLEDGE_BASE",
Review Comment:
I have added it in, but I did not realize that the "external sources" feature was only recently added in botocore. I tested it by launching breeze and running `pip install boto3==1.34.90` then running the dag. So it does work IF the user installs a version of boto which might cause issues with aiobotocore.
The latest version of aiobotocore constrains botocore<1.34.70,>=1.34.41.
I'm not entirely sure what the best way to handle this would be. Very open to suggestions and advice on this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1598853590
##########
tests/system/providers/amazon/aws/example_bedrock_knowledge_base.py:
##########
@@ -480,6 +483,24 @@ def delete_opensearch_policies(collection_name: str):
)
# [END howto_sensor_bedrock_ingest_data]
+ # [START howto_operator_bedrock_retrieve_and_generate]
+ retrieve_and_generate = BedrockRaGOperator(
+ task_id="retrieve_and_generate",
+ input="Who was the CEO of Amazon on 2022?",
+ source_type="KNOWLEDGE_BASE",
Review Comment:
Simpler code, but less explicit maybe. Plus it'd help to log the version of boto that folks have so that they know where they're at "<version X> required and found <version Y>". So maybe your approach is better that way.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1596046080
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -664,3 +669,198 @@ def execute(self, context: Context) -> str:
)
return ingestion_job_id
+
+
+class BedrockRaGOperator(AwsBaseOperator[BedrockAgentRuntimeHook]):
+ """
+ Query a knowledge base and generate responses based on the retrieved results with sources citations.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockRaGOperator`
+
+ :param input: The query to be made to the knowledge base. (templated)
+ :param source_type: The type of resource that is queried by the request. (templated)
+ Must be one of 'KNOWLEDGE_BASE' or 'EXTERNAL_SOURCES', and the appropriate config values must also be provided.
+ If set to 'KNOWLEDGE_BASE' then `knowledge_base_id` must be provided, and `vector_search_config` may be.
+ If set to `EXTERNAL_SOURCES` then `sources` must also be provided.
+ :param model_arn: The ARN of the foundation model used to generate a response. (templated)
+ :param prompt_template: The template for the prompt that's sent to the model for response generation.
+ You can include prompt placeholders, which become replaced before the prompt is sent to the model
+ to provide instructions and context to the model. In addition, you can include XML tags to delineate
+ meaningful sections of the prompt template. (templated)
+ :param knowledge_base_id: The unique identifier of the knowledge base that is queried. (templated)
+ Can only be specified if source_type='KNOWLEDGE_BASE'.
+ :param vector_search_config: How the results from the vector search should be returned. (templated)
+ Can only be specified if source_type='KNOWLEDGE_BASE'.
+ For more information, see https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-config.html.
+ :param sources: The documents used as reference for the response. (templated)
+ Can only be specified if source_type='EXTERNAL_SOURCES'
+ :param rag_kwargs: Additional keyword arguments to pass to the API call. (templated)
+ """
+
+ aws_hook_class = BedrockAgentRuntimeHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "input",
+ "source_type",
+ "model_arn",
+ "prompt_template",
+ "knowledge_base_id",
+ "vector_search_config",
+ "sources",
+ "rag_kwargs",
+ )
+
+ def __init__(
+ self,
+ input: str,
+ source_type: str,
+ model_arn: str,
+ prompt_template: str | None = None,
+ knowledge_base_id: str | None = None,
+ vector_search_config: dict[str, Any] | None = None,
+ sources: list[dict[str, Any]] | None = None,
+ rag_kwargs: dict[str, Any] | None = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.input = input
+ self.prompt_template = prompt_template
+ self.source_type = source_type.upper()
+ self.knowledge_base_id = knowledge_base_id
+ self.model_arn = model_arn
+ self.vector_search_config = vector_search_config
+ self.sources = sources
+ self.rag_kwargs = rag_kwargs or {}
+
+ def validate_inputs(self):
+ if self.source_type == "KNOWLEDGE_BASE":
+ if self.knowledge_base_id is None:
+ raise AttributeError(
+ "If `source_type` is set to 'KNOWLEDGE_BASE' then `knowledge_base_id` must be provided."
+ )
+ if self.sources is not None:
+ raise AttributeError(
+ "`sources` can not be used when `source_type` is set to 'KNOWLEDGE_BASE'."
+ )
+ elif self.source_type == "EXTERNAL_SOURCES":
+ if not self.sources is not None:
+ raise AttributeError(
+ "If `source_type` is set to `EXTERNAL_SOURCES` then `sources` must also be provided."
+ )
+ if self.vector_search_config or self.knowledge_base_id:
+ raise AttributeError(
+ "`vector_search_config` and `knowledge_base_id` can not be used "
+ "when `source_type` is set to `EXTERNAL_SOURCES`"
+ )
+ else:
+ raise AttributeError(
+ "`source_type` must be one of 'KNOWLEDGE_BASE' or 'EXTERNAL_SOURCES', "
+ "and the appropriate config values must also be provided."
+ )
+
+ def build_rag_config(self) -> dict[str, Any]:
+ result: dict[str, Any] = {}
+ base_config: dict[str, Any] = {
+ "modelArn": self.model_arn,
+ }
+
+ if self.prompt_template:
+ base_config["generationConfiguration"] = {
+ "promptTemplate": {"textPromptTemplate": self.prompt_template}
+ }
+
+ if self.source_type == "KNOWLEDGE_BASE":
+ if self.vector_search_config:
+ base_config["retrievalConfiguration"] = {
+ "vectorSearchConfiguration": self.vector_search_config
+ }
+
+ result = {
+ "type": self.source_type,
+ "knowledgeBaseConfiguration": {
+ **base_config,
+ "knowledgeBaseId": self.knowledge_base_id,
+ },
+ }
+
+ if self.source_type == "EXTERNAL_SOURCES":
+ result = {
+ "type": self.source_type,
+ "externalSourcesConfiguration": {**base_config, "sources": self.sources},
+ }
+ return result
+
+ def execute(self, context: Context) -> Any:
+ self.validate_inputs()
+
+ result = self.hook.conn.retrieve_and_generate(
+ input={"text": self.input},
+ retrieveAndGenerateConfiguration=self.build_rag_config(),
+ **self.rag_kwargs,
+ )
+
+ self.log.info(
+ "\nPrompt: %s\nResponse: %s\nCitations: %s",
+ self.input,
+ result["output"]["text"],
+ result["citations"],
+ )
+ return result
+
+
+class BedrockRetrieveOperator(AwsBaseOperator[BedrockAgentRuntimeHook]):
+ """
+ Query a knowledge base and retrieve results with source citations.
Review Comment:
`Retrieve` returns only the citations without a generated answer. `RaG` will return the citations and a pretty LLM-generated answer.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1596045184
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -664,3 +669,198 @@ def execute(self, context: Context) -> str:
)
return ingestion_job_id
+
+
+class BedrockRaGOperator(AwsBaseOperator[BedrockAgentRuntimeHook]):
+ """
+ Query a knowledge base and generate responses based on the retrieved results with sources citations.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockRaGOperator`
+
+ :param input: The query to be made to the knowledge base. (templated)
+ :param source_type: The type of resource that is queried by the request. (templated)
+ Must be one of 'KNOWLEDGE_BASE' or 'EXTERNAL_SOURCES', and the appropriate config values must also be provided.
+ If set to 'KNOWLEDGE_BASE' then `knowledge_base_id` must be provided, and `vector_search_config` may be.
+ If set to `EXTERNAL_SOURCES` then `sources` must also be provided.
+ :param model_arn: The ARN of the foundation model used to generate a response. (templated)
+ :param prompt_template: The template for the prompt that's sent to the model for response generation.
+ You can include prompt placeholders, which become replaced before the prompt is sent to the model
Review Comment:
Heh. that's copypasta from their docs, but I agree. I'll make the change here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1598821701
##########
tests/system/providers/amazon/aws/example_bedrock_knowledge_base.py:
##########
@@ -480,6 +483,24 @@ def delete_opensearch_policies(collection_name: str):
)
# [END howto_sensor_bedrock_ingest_data]
+ # [START howto_operator_bedrock_retrieve_and_generate]
+ retrieve_and_generate = BedrockRaGOperator(
+ task_id="retrieve_and_generate",
+ input="Who was the CEO of Amazon on 2022?",
+ source_type="KNOWLEDGE_BASE",
Review Comment:
Yeah, I agree with this. Though, I wouldn't wrap in a version check, you can just handle the exception you get back from the boto sdk right?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1596050868
##########
docs/apache-airflow-providers-amazon/operators/bedrock.rst:
##########
@@ -174,6 +177,40 @@ To add data from an Amazon S3 bucket into an Amazon Bedrock Data Source, you can
:start-after: [START howto_operator_bedrock_ingest_data]
:end-before: [END howto_operator_bedrock_ingest_data]
+.. _howto/operator:BedrockRetrieveOperator:
+
+Amazon Bedrock Retrieve
+=======================
+
+To query a knowledge base, you can use :class:`~airflow.providers.amazon.aws.operators.bedrock.BedrockRaGOperator`.
Review Comment:
> This should be the BedrockRetrieveOperator, right?
You are right
> I also don't fully understand the difference
I'll add some more context tot he docs and see if that helps.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1596061444
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -664,3 +669,198 @@ def execute(self, context: Context) -> str:
)
return ingestion_job_id
+
+
+class BedrockRaGOperator(AwsBaseOperator[BedrockAgentRuntimeHook]):
+ """
+ Query a knowledge base and generate responses based on the retrieved results with sources citations.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockRaGOperator`
+
+ :param input: The query to be made to the knowledge base. (templated)
+ :param source_type: The type of resource that is queried by the request. (templated)
+ Must be one of 'KNOWLEDGE_BASE' or 'EXTERNAL_SOURCES', and the appropriate config values must also be provided.
+ If set to 'KNOWLEDGE_BASE' then `knowledge_base_id` must be provided, and `vector_search_config` may be.
+ If set to `EXTERNAL_SOURCES` then `sources` must also be provided.
+ :param model_arn: The ARN of the foundation model used to generate a response. (templated)
+ :param prompt_template: The template for the prompt that's sent to the model for response generation.
+ You can include prompt placeholders, which become replaced before the prompt is sent to the model
+ to provide instructions and context to the model. In addition, you can include XML tags to delineate
+ meaningful sections of the prompt template. (templated)
+ :param knowledge_base_id: The unique identifier of the knowledge base that is queried. (templated)
+ Can only be specified if source_type='KNOWLEDGE_BASE'.
+ :param vector_search_config: How the results from the vector search should be returned. (templated)
+ Can only be specified if source_type='KNOWLEDGE_BASE'.
+ For more information, see https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-config.html.
+ :param sources: The documents used as reference for the response. (templated)
+ Can only be specified if source_type='EXTERNAL_SOURCES'
+ :param rag_kwargs: Additional keyword arguments to pass to the API call. (templated)
+ """
+
+ aws_hook_class = BedrockAgentRuntimeHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "input",
+ "source_type",
+ "model_arn",
+ "prompt_template",
+ "knowledge_base_id",
+ "vector_search_config",
+ "sources",
+ "rag_kwargs",
+ )
+
+ def __init__(
+ self,
+ input: str,
+ source_type: str,
+ model_arn: str,
+ prompt_template: str | None = None,
+ knowledge_base_id: str | None = None,
+ vector_search_config: dict[str, Any] | None = None,
+ sources: list[dict[str, Any]] | None = None,
+ rag_kwargs: dict[str, Any] | None = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.input = input
+ self.prompt_template = prompt_template
+ self.source_type = source_type.upper()
+ self.knowledge_base_id = knowledge_base_id
+ self.model_arn = model_arn
+ self.vector_search_config = vector_search_config
+ self.sources = sources
+ self.rag_kwargs = rag_kwargs or {}
+
+ def validate_inputs(self):
+ if self.source_type == "KNOWLEDGE_BASE":
+ if self.knowledge_base_id is None:
+ raise AttributeError(
+ "If `source_type` is set to 'KNOWLEDGE_BASE' then `knowledge_base_id` must be provided."
+ )
+ if self.sources is not None:
+ raise AttributeError(
+ "`sources` can not be used when `source_type` is set to 'KNOWLEDGE_BASE'."
+ )
+ elif self.source_type == "EXTERNAL_SOURCES":
+ if not self.sources is not None:
+ raise AttributeError(
+ "If `source_type` is set to `EXTERNAL_SOURCES` then `sources` must also be provided."
+ )
+ if self.vector_search_config or self.knowledge_base_id:
+ raise AttributeError(
+ "`vector_search_config` and `knowledge_base_id` can not be used "
+ "when `source_type` is set to `EXTERNAL_SOURCES`"
+ )
+ else:
+ raise AttributeError(
+ "`source_type` must be one of 'KNOWLEDGE_BASE' or 'EXTERNAL_SOURCES', "
+ "and the appropriate config values must also be provided."
+ )
+
+ def build_rag_config(self) -> dict[str, Any]:
+ result: dict[str, Any] = {}
+ base_config: dict[str, Any] = {
+ "modelArn": self.model_arn,
+ }
+
+ if self.prompt_template:
+ base_config["generationConfiguration"] = {
+ "promptTemplate": {"textPromptTemplate": self.prompt_template}
+ }
+
+ if self.source_type == "KNOWLEDGE_BASE":
+ if self.vector_search_config:
+ base_config["retrievalConfiguration"] = {
+ "vectorSearchConfiguration": self.vector_search_config
+ }
+
+ result = {
+ "type": self.source_type,
+ "knowledgeBaseConfiguration": {
+ **base_config,
+ "knowledgeBaseId": self.knowledge_base_id,
+ },
+ }
+
+ if self.source_type == "EXTERNAL_SOURCES":
+ result = {
+ "type": self.source_type,
+ "externalSourcesConfiguration": {**base_config, "sources": self.sources},
+ }
+ return result
+
+ def execute(self, context: Context) -> Any:
+ self.validate_inputs()
+
+ result = self.hook.conn.retrieve_and_generate(
+ input={"text": self.input},
+ retrieveAndGenerateConfiguration=self.build_rag_config(),
+ **self.rag_kwargs,
+ )
+
+ self.log.info(
+ "\nPrompt: %s\nResponse: %s\nCitations: %s",
+ self.input,
+ result["output"]["text"],
+ result["citations"],
+ )
+ return result
+
+
+class BedrockRetrieveOperator(AwsBaseOperator[BedrockAgentRuntimeHook]):
+ """
+ Query a knowledge base and retrieve results with source citations.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockRetrieveOperator`
+
+ :param retrieval_query: The query to be made to the knowledge base. (templated)
+ :param knowledge_base_id: The unique identifier of the knowledge base that is queried. (templated)
+ Can only be specified if source_type='KNOWLEDGE_BASE'.
Review Comment:
Yup, good eye. The `Retrieve` API only supports knowledge bases, so it's not needed here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1598804403
##########
tests/system/providers/amazon/aws/example_bedrock_knowledge_base.py:
##########
@@ -480,6 +483,24 @@ def delete_opensearch_policies(collection_name: str):
)
# [END howto_sensor_bedrock_ingest_data]
+ # [START howto_operator_bedrock_retrieve_and_generate]
+ retrieve_and_generate = BedrockRaGOperator(
+ task_id="retrieve_and_generate",
+ input="Who was the CEO of Amazon on 2022?",
+ source_type="KNOWLEDGE_BASE",
Review Comment:
While adding this I found out that “external sources” was only added to the API 3 weeks ago, meaning it would need botocore min version 1.34.103 current dependency is botocore>=1.33.0 which leaves us with an awkward situation. We can bump minver on botocore but forcing a 3-week-old version feels aggressive.
We don't usually do this, but I think in this case maybe wrapping it in a botocore version check and adding a note inthe docs may be the right call?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1596065112
##########
tests/providers/amazon/aws/operators/test_bedrock.py:
##########
@@ -346,3 +347,168 @@ def test_id_returned(self, mock_conn):
result = self.operator.execute({})
assert result == self.INGESTION_JOB_ID
+
+
+class TestBedrockRaGOperator:
+ VECTOR_SEARCH_CONFIG = {"filter": {"equals": {"key": "some key", "value": "some value"}}}
+ KNOWLEDGE_BASE_ID = "knowledge_base_id"
+ SOURCES = [{"sourceType": "S3", "s3Location": "bucket"}]
+ MODEL_ARN = "model arn"
+
+ @pytest.mark.parametrize(
+ "source_type, vector_search_config, knowledge_base_id, sources, expect_success",
+ [
+ pytest.param(
+ "invalid_source_type",
+ None,
+ None,
+ None,
+ False,
+ id="invalid_source_type",
+ ),
+ pytest.param(
+ "KNOWLEDGE_BASE",
+ VECTOR_SEARCH_CONFIG,
+ None,
+ None,
+ False,
+ id="KNOWLEDGE_BASE_without_knowledge_base_id_fails",
+ ),
+ pytest.param(
+ "KNOWLEDGE_BASE",
+ VECTOR_SEARCH_CONFIG,
+ KNOWLEDGE_BASE_ID,
+ SOURCES,
+ False,
+ id="KNOWLEDGE_BASE_with_sources_fails",
+ ),
+ pytest.param(
+ "KNOWLEDGE_BASE",
+ VECTOR_SEARCH_CONFIG,
+ KNOWLEDGE_BASE_ID,
+ None,
+ True,
+ id="KNOWLEDGE_BASE_passes",
+ ),
Review Comment:
Sure. vector config is optional with a knowledge base, but forbidden when using external. I can add that check.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1596061831
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -664,3 +669,198 @@ def execute(self, context: Context) -> str:
)
return ingestion_job_id
+
+
+class BedrockRaGOperator(AwsBaseOperator[BedrockAgentRuntimeHook]):
+ """
+ Query a knowledge base and generate responses based on the retrieved results with sources citations.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockRaGOperator`
+
+ :param input: The query to be made to the knowledge base. (templated)
+ :param source_type: The type of resource that is queried by the request. (templated)
+ Must be one of 'KNOWLEDGE_BASE' or 'EXTERNAL_SOURCES', and the appropriate config values must also be provided.
+ If set to 'KNOWLEDGE_BASE' then `knowledge_base_id` must be provided, and `vector_search_config` may be.
+ If set to `EXTERNAL_SOURCES` then `sources` must also be provided.
+ :param model_arn: The ARN of the foundation model used to generate a response. (templated)
+ :param prompt_template: The template for the prompt that's sent to the model for response generation.
+ You can include prompt placeholders, which become replaced before the prompt is sent to the model
+ to provide instructions and context to the model. In addition, you can include XML tags to delineate
+ meaningful sections of the prompt template. (templated)
+ :param knowledge_base_id: The unique identifier of the knowledge base that is queried. (templated)
+ Can only be specified if source_type='KNOWLEDGE_BASE'.
+ :param vector_search_config: How the results from the vector search should be returned. (templated)
+ Can only be specified if source_type='KNOWLEDGE_BASE'.
+ For more information, see https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-config.html.
+ :param sources: The documents used as reference for the response. (templated)
+ Can only be specified if source_type='EXTERNAL_SOURCES'
+ :param rag_kwargs: Additional keyword arguments to pass to the API call. (templated)
+ """
+
+ aws_hook_class = BedrockAgentRuntimeHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "input",
+ "source_type",
+ "model_arn",
+ "prompt_template",
+ "knowledge_base_id",
+ "vector_search_config",
+ "sources",
+ "rag_kwargs",
+ )
+
+ def __init__(
+ self,
+ input: str,
+ source_type: str,
+ model_arn: str,
+ prompt_template: str | None = None,
+ knowledge_base_id: str | None = None,
+ vector_search_config: dict[str, Any] | None = None,
+ sources: list[dict[str, Any]] | None = None,
+ rag_kwargs: dict[str, Any] | None = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.input = input
+ self.prompt_template = prompt_template
+ self.source_type = source_type.upper()
+ self.knowledge_base_id = knowledge_base_id
+ self.model_arn = model_arn
+ self.vector_search_config = vector_search_config
+ self.sources = sources
+ self.rag_kwargs = rag_kwargs or {}
+
+ def validate_inputs(self):
+ if self.source_type == "KNOWLEDGE_BASE":
+ if self.knowledge_base_id is None:
+ raise AttributeError(
+ "If `source_type` is set to 'KNOWLEDGE_BASE' then `knowledge_base_id` must be provided."
+ )
+ if self.sources is not None:
+ raise AttributeError(
+ "`sources` can not be used when `source_type` is set to 'KNOWLEDGE_BASE'."
+ )
+ elif self.source_type == "EXTERNAL_SOURCES":
+ if not self.sources is not None:
+ raise AttributeError(
+ "If `source_type` is set to `EXTERNAL_SOURCES` then `sources` must also be provided."
+ )
+ if self.vector_search_config or self.knowledge_base_id:
+ raise AttributeError(
+ "`vector_search_config` and `knowledge_base_id` can not be used "
+ "when `source_type` is set to `EXTERNAL_SOURCES`"
+ )
+ else:
+ raise AttributeError(
+ "`source_type` must be one of 'KNOWLEDGE_BASE' or 'EXTERNAL_SOURCES', "
+ "and the appropriate config values must also be provided."
+ )
+
+ def build_rag_config(self) -> dict[str, Any]:
+ result: dict[str, Any] = {}
+ base_config: dict[str, Any] = {
+ "modelArn": self.model_arn,
+ }
+
+ if self.prompt_template:
+ base_config["generationConfiguration"] = {
+ "promptTemplate": {"textPromptTemplate": self.prompt_template}
+ }
+
+ if self.source_type == "KNOWLEDGE_BASE":
+ if self.vector_search_config:
+ base_config["retrievalConfiguration"] = {
+ "vectorSearchConfiguration": self.vector_search_config
+ }
+
+ result = {
+ "type": self.source_type,
+ "knowledgeBaseConfiguration": {
+ **base_config,
+ "knowledgeBaseId": self.knowledge_base_id,
+ },
+ }
+
+ if self.source_type == "EXTERNAL_SOURCES":
+ result = {
+ "type": self.source_type,
+ "externalSourcesConfiguration": {**base_config, "sources": self.sources},
+ }
+ return result
+
+ def execute(self, context: Context) -> Any:
+ self.validate_inputs()
+
+ result = self.hook.conn.retrieve_and_generate(
+ input={"text": self.input},
+ retrieveAndGenerateConfiguration=self.build_rag_config(),
+ **self.rag_kwargs,
+ )
+
+ self.log.info(
+ "\nPrompt: %s\nResponse: %s\nCitations: %s",
+ self.input,
+ result["output"]["text"],
+ result["citations"],
+ )
+ return result
+
+
+class BedrockRetrieveOperator(AwsBaseOperator[BedrockAgentRuntimeHook]):
+ """
+ Query a knowledge base and retrieve results with source citations.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockRetrieveOperator`
+
+ :param retrieval_query: The query to be made to the knowledge base. (templated)
+ :param knowledge_base_id: The unique identifier of the knowledge base that is queried. (templated)
+ Can only be specified if source_type='KNOWLEDGE_BASE'.
Review Comment:
OH! but I borked the docstring! I see what you meant. :facepalm:
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on PR #39500:
URL: https://github.com/apache/airflow/pull/39500#issuecomment-2113133539
Ran it using the system test trigger script and it passes with the external sources getting skipped:
```
INFO [tests.system.utils] EXECUTION SUMMARY:
╒═════════════════════════════════════════════════╤══════════╕
│ Task ID │ Status │
╞═════════════════════════════════════════════════╪══════════╡
│ variable_fetcher │ success │
├─────────────────────────────────────────────────┼──────────┤
│ create_bucket │ success │
├─────────────────────────────────────────────────┼──────────┤
│ create_opensearch_policies │ success │
├─────────────────────────────────────────────────┼──────────┤
│ create_collection │ success │
├─────────────────────────────────────────────────┼──────────┤
│ await_collection │ success │
├─────────────────────────────────────────────────┼──────────┤
│ get_collection_arn │ success │
├─────────────────────────────────────────────────┼──────────┤
│ create_vector_index │ success │
├─────────────────────────────────────────────────┼──────────┤
│ copy_data_to_s3 │ success │
├─────────────────────────────────────────────────┼──────────┤
│ create_knowledge_base │ success │
├─────────────────────────────────────────────────┼──────────┤
│ await_knowledge_base │ success │
├─────────────────────────────────────────────────┼──────────┤
│ create_data_source │ success │
├─────────────────────────────────────────────────┼──────────┤
│ ingest_data │ success │
├─────────────────────────────────────────────────┼──────────┤
│ await_ingest │ success │
├─────────────────────────────────────────────────┼──────────┤
│ knowledge_base_rag │ success │
├─────────────────────────────────────────────────┼──────────┤
│ external_sources_rag_group.external_sources_rag │ skipped │
├─────────────────────────────────────────────────┼──────────┤
│ external_sources_rag_group.run_or_skip │ success │
├─────────────────────────────────────────────────┼──────────┤
│ external_sources_rag_group.end_workflow │ success │
├─────────────────────────────────────────────────┼──────────┤
│ retrieve │ success │
├─────────────────────────────────────────────────┼──────────┤
│ delete_data_source │ success │
├─────────────────────────────────────────────────┼──────────┤
│ delete_knowledge_base │ success │
├─────────────────────────────────────────────────┼──────────┤
│ delete_vector_index │ success │
├─────────────────────────────────────────────────┼──────────┤
│ delete_opensearch_policies │ success │
├─────────────────────────────────────────────────┼──────────┤
│ delete_collection │ success │
├─────────────────────────────────────────────────┼──────────┤
│ delete_bucket │ success │
├─────────────────────────────────────────────────┼──────────┤
│ watcher │ skipped │
╘═════════════════════════════════════════════════╧══════════╛
```
```
======================== 1 passed in 423.02s (0:07:03) =========================
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1596070383
##########
tests/system/providers/amazon/aws/example_bedrock_knowledge_base.py:
##########
@@ -480,6 +483,24 @@ def delete_opensearch_policies(collection_name: str):
)
# [END howto_sensor_bedrock_ingest_data]
+ # [START howto_operator_bedrock_retrieve_and_generate]
+ retrieve_and_generate = BedrockRaGOperator(
+ task_id="retrieve_and_generate",
+ input="Who was the CEO of Amazon on 2022?",
+ source_type="KNOWLEDGE_BASE",
Review Comment:
Sure, I should be able to add that with the s3 data I'm already uploading as part of the test I think... I'll have to experiment a little, but I think I can do it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1596032939
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -664,3 +669,198 @@ def execute(self, context: Context) -> str:
)
return ingestion_job_id
+
+
+class BedrockRaGOperator(AwsBaseOperator[BedrockAgentRuntimeHook]):
+ """
+ Query a knowledge base and generate responses based on the retrieved results with sources citations.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockRaGOperator`
+
+ :param input: The query to be made to the knowledge base. (templated)
+ :param source_type: The type of resource that is queried by the request. (templated)
+ Must be one of 'KNOWLEDGE_BASE' or 'EXTERNAL_SOURCES', and the appropriate config values must also be provided.
+ If set to 'KNOWLEDGE_BASE' then `knowledge_base_id` must be provided, and `vector_search_config` may be.
+ If set to `EXTERNAL_SOURCES` then `sources` must also be provided.
+ :param model_arn: The ARN of the foundation model used to generate a response. (templated)
+ :param prompt_template: The template for the prompt that's sent to the model for response generation.
+ You can include prompt placeholders, which become replaced before the prompt is sent to the model
Review Comment:
```suggestion
You can include prompt placeholders, which are replaced before the prompt is sent to the model
```
##########
docs/apache-airflow-providers-amazon/operators/bedrock.rst:
##########
@@ -174,6 +177,40 @@ To add data from an Amazon S3 bucket into an Amazon Bedrock Data Source, you can
:start-after: [START howto_operator_bedrock_ingest_data]
:end-before: [END howto_operator_bedrock_ingest_data]
+.. _howto/operator:BedrockRetrieveOperator:
+
+Amazon Bedrock Retrieve
+=======================
+
+To query a knowledge base, you can use :class:`~airflow.providers.amazon.aws.operators.bedrock.BedrockRaGOperator`.
Review Comment:
This should be the `BedrockRetrieveOperator`, right? Or maybe the one below? Right now RAG is in both. I also don't fully understand the difference between the two operators. Maybe just a sentence or two more to tell the two apart? The only difference I see in the description is one says "includes" the other says "cites"
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -664,3 +669,198 @@ def execute(self, context: Context) -> str:
)
return ingestion_job_id
+
+
+class BedrockRaGOperator(AwsBaseOperator[BedrockAgentRuntimeHook]):
+ """
+ Query a knowledge base and generate responses based on the retrieved results with sources citations.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockRaGOperator`
+
+ :param input: The query to be made to the knowledge base. (templated)
+ :param source_type: The type of resource that is queried by the request. (templated)
+ Must be one of 'KNOWLEDGE_BASE' or 'EXTERNAL_SOURCES', and the appropriate config values must also be provided.
+ If set to 'KNOWLEDGE_BASE' then `knowledge_base_id` must be provided, and `vector_search_config` may be.
+ If set to `EXTERNAL_SOURCES` then `sources` must also be provided.
+ :param model_arn: The ARN of the foundation model used to generate a response. (templated)
+ :param prompt_template: The template for the prompt that's sent to the model for response generation.
+ You can include prompt placeholders, which become replaced before the prompt is sent to the model
+ to provide instructions and context to the model. In addition, you can include XML tags to delineate
+ meaningful sections of the prompt template. (templated)
+ :param knowledge_base_id: The unique identifier of the knowledge base that is queried. (templated)
+ Can only be specified if source_type='KNOWLEDGE_BASE'.
+ :param vector_search_config: How the results from the vector search should be returned. (templated)
+ Can only be specified if source_type='KNOWLEDGE_BASE'.
+ For more information, see https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-config.html.
+ :param sources: The documents used as reference for the response. (templated)
+ Can only be specified if source_type='EXTERNAL_SOURCES'
+ :param rag_kwargs: Additional keyword arguments to pass to the API call. (templated)
+ """
+
+ aws_hook_class = BedrockAgentRuntimeHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "input",
+ "source_type",
+ "model_arn",
+ "prompt_template",
+ "knowledge_base_id",
+ "vector_search_config",
+ "sources",
+ "rag_kwargs",
+ )
+
+ def __init__(
+ self,
+ input: str,
+ source_type: str,
+ model_arn: str,
+ prompt_template: str | None = None,
+ knowledge_base_id: str | None = None,
+ vector_search_config: dict[str, Any] | None = None,
+ sources: list[dict[str, Any]] | None = None,
+ rag_kwargs: dict[str, Any] | None = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.input = input
+ self.prompt_template = prompt_template
+ self.source_type = source_type.upper()
+ self.knowledge_base_id = knowledge_base_id
+ self.model_arn = model_arn
+ self.vector_search_config = vector_search_config
+ self.sources = sources
+ self.rag_kwargs = rag_kwargs or {}
+
+ def validate_inputs(self):
+ if self.source_type == "KNOWLEDGE_BASE":
+ if self.knowledge_base_id is None:
+ raise AttributeError(
+ "If `source_type` is set to 'KNOWLEDGE_BASE' then `knowledge_base_id` must be provided."
+ )
+ if self.sources is not None:
+ raise AttributeError(
+ "`sources` can not be used when `source_type` is set to 'KNOWLEDGE_BASE'."
+ )
+ elif self.source_type == "EXTERNAL_SOURCES":
+ if not self.sources is not None:
+ raise AttributeError(
+ "If `source_type` is set to `EXTERNAL_SOURCES` then `sources` must also be provided."
+ )
+ if self.vector_search_config or self.knowledge_base_id:
+ raise AttributeError(
+ "`vector_search_config` and `knowledge_base_id` can not be used "
+ "when `source_type` is set to `EXTERNAL_SOURCES`"
+ )
+ else:
+ raise AttributeError(
+ "`source_type` must be one of 'KNOWLEDGE_BASE' or 'EXTERNAL_SOURCES', "
+ "and the appropriate config values must also be provided."
+ )
+
+ def build_rag_config(self) -> dict[str, Any]:
+ result: dict[str, Any] = {}
+ base_config: dict[str, Any] = {
+ "modelArn": self.model_arn,
+ }
+
+ if self.prompt_template:
+ base_config["generationConfiguration"] = {
+ "promptTemplate": {"textPromptTemplate": self.prompt_template}
+ }
+
+ if self.source_type == "KNOWLEDGE_BASE":
+ if self.vector_search_config:
+ base_config["retrievalConfiguration"] = {
+ "vectorSearchConfiguration": self.vector_search_config
+ }
+
+ result = {
+ "type": self.source_type,
+ "knowledgeBaseConfiguration": {
+ **base_config,
+ "knowledgeBaseId": self.knowledge_base_id,
+ },
+ }
+
+ if self.source_type == "EXTERNAL_SOURCES":
+ result = {
+ "type": self.source_type,
+ "externalSourcesConfiguration": {**base_config, "sources": self.sources},
+ }
+ return result
+
+ def execute(self, context: Context) -> Any:
+ self.validate_inputs()
+
+ result = self.hook.conn.retrieve_and_generate(
+ input={"text": self.input},
+ retrieveAndGenerateConfiguration=self.build_rag_config(),
+ **self.rag_kwargs,
+ )
+
+ self.log.info(
+ "\nPrompt: %s\nResponse: %s\nCitations: %s",
+ self.input,
+ result["output"]["text"],
+ result["citations"],
+ )
+ return result
+
+
+class BedrockRetrieveOperator(AwsBaseOperator[BedrockAgentRuntimeHook]):
+ """
+ Query a knowledge base and retrieve results with source citations.
Review Comment:
Does this one also retrieve citations? I thought it was just the RAG operator that did that?
##########
tests/providers/amazon/aws/operators/test_bedrock.py:
##########
@@ -346,3 +347,168 @@ def test_id_returned(self, mock_conn):
result = self.operator.execute({})
assert result == self.INGESTION_JOB_ID
+
+
+class TestBedrockRaGOperator:
+ VECTOR_SEARCH_CONFIG = {"filter": {"equals": {"key": "some key", "value": "some value"}}}
+ KNOWLEDGE_BASE_ID = "knowledge_base_id"
+ SOURCES = [{"sourceType": "S3", "s3Location": "bucket"}]
+ MODEL_ARN = "model arn"
+
+ @pytest.mark.parametrize(
+ "source_type, vector_search_config, knowledge_base_id, sources, expect_success",
+ [
+ pytest.param(
+ "invalid_source_type",
+ None,
+ None,
+ None,
+ False,
+ id="invalid_source_type",
+ ),
+ pytest.param(
+ "KNOWLEDGE_BASE",
+ VECTOR_SEARCH_CONFIG,
+ None,
+ None,
+ False,
+ id="KNOWLEDGE_BASE_without_knowledge_base_id_fails",
+ ),
+ pytest.param(
+ "KNOWLEDGE_BASE",
+ VECTOR_SEARCH_CONFIG,
+ KNOWLEDGE_BASE_ID,
+ SOURCES,
+ False,
+ id="KNOWLEDGE_BASE_with_sources_fails",
+ ),
+ pytest.param(
+ "KNOWLEDGE_BASE",
+ VECTOR_SEARCH_CONFIG,
+ KNOWLEDGE_BASE_ID,
+ None,
+ True,
+ id="KNOWLEDGE_BASE_passes",
+ ),
Review Comment:
Should another test case with just a knowledge_base_id (without vector config) also pass?
##########
airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -664,3 +669,198 @@ def execute(self, context: Context) -> str:
)
return ingestion_job_id
+
+
+class BedrockRaGOperator(AwsBaseOperator[BedrockAgentRuntimeHook]):
+ """
+ Query a knowledge base and generate responses based on the retrieved results with sources citations.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockRaGOperator`
+
+ :param input: The query to be made to the knowledge base. (templated)
+ :param source_type: The type of resource that is queried by the request. (templated)
+ Must be one of 'KNOWLEDGE_BASE' or 'EXTERNAL_SOURCES', and the appropriate config values must also be provided.
+ If set to 'KNOWLEDGE_BASE' then `knowledge_base_id` must be provided, and `vector_search_config` may be.
+ If set to `EXTERNAL_SOURCES` then `sources` must also be provided.
+ :param model_arn: The ARN of the foundation model used to generate a response. (templated)
+ :param prompt_template: The template for the prompt that's sent to the model for response generation.
+ You can include prompt placeholders, which become replaced before the prompt is sent to the model
+ to provide instructions and context to the model. In addition, you can include XML tags to delineate
+ meaningful sections of the prompt template. (templated)
+ :param knowledge_base_id: The unique identifier of the knowledge base that is queried. (templated)
+ Can only be specified if source_type='KNOWLEDGE_BASE'.
+ :param vector_search_config: How the results from the vector search should be returned. (templated)
+ Can only be specified if source_type='KNOWLEDGE_BASE'.
+ For more information, see https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-config.html.
+ :param sources: The documents used as reference for the response. (templated)
+ Can only be specified if source_type='EXTERNAL_SOURCES'
+ :param rag_kwargs: Additional keyword arguments to pass to the API call. (templated)
+ """
+
+ aws_hook_class = BedrockAgentRuntimeHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "input",
+ "source_type",
+ "model_arn",
+ "prompt_template",
+ "knowledge_base_id",
+ "vector_search_config",
+ "sources",
+ "rag_kwargs",
+ )
+
+ def __init__(
+ self,
+ input: str,
+ source_type: str,
+ model_arn: str,
+ prompt_template: str | None = None,
+ knowledge_base_id: str | None = None,
+ vector_search_config: dict[str, Any] | None = None,
+ sources: list[dict[str, Any]] | None = None,
+ rag_kwargs: dict[str, Any] | None = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.input = input
+ self.prompt_template = prompt_template
+ self.source_type = source_type.upper()
+ self.knowledge_base_id = knowledge_base_id
+ self.model_arn = model_arn
+ self.vector_search_config = vector_search_config
+ self.sources = sources
+ self.rag_kwargs = rag_kwargs or {}
+
+ def validate_inputs(self):
+ if self.source_type == "KNOWLEDGE_BASE":
+ if self.knowledge_base_id is None:
+ raise AttributeError(
+ "If `source_type` is set to 'KNOWLEDGE_BASE' then `knowledge_base_id` must be provided."
+ )
+ if self.sources is not None:
+ raise AttributeError(
+ "`sources` can not be used when `source_type` is set to 'KNOWLEDGE_BASE'."
+ )
+ elif self.source_type == "EXTERNAL_SOURCES":
+ if not self.sources is not None:
+ raise AttributeError(
+ "If `source_type` is set to `EXTERNAL_SOURCES` then `sources` must also be provided."
+ )
+ if self.vector_search_config or self.knowledge_base_id:
+ raise AttributeError(
+ "`vector_search_config` and `knowledge_base_id` can not be used "
+ "when `source_type` is set to `EXTERNAL_SOURCES`"
+ )
+ else:
+ raise AttributeError(
+ "`source_type` must be one of 'KNOWLEDGE_BASE' or 'EXTERNAL_SOURCES', "
+ "and the appropriate config values must also be provided."
+ )
+
+ def build_rag_config(self) -> dict[str, Any]:
+ result: dict[str, Any] = {}
+ base_config: dict[str, Any] = {
+ "modelArn": self.model_arn,
+ }
+
+ if self.prompt_template:
+ base_config["generationConfiguration"] = {
+ "promptTemplate": {"textPromptTemplate": self.prompt_template}
+ }
+
+ if self.source_type == "KNOWLEDGE_BASE":
+ if self.vector_search_config:
+ base_config["retrievalConfiguration"] = {
+ "vectorSearchConfiguration": self.vector_search_config
+ }
+
+ result = {
+ "type": self.source_type,
+ "knowledgeBaseConfiguration": {
+ **base_config,
+ "knowledgeBaseId": self.knowledge_base_id,
+ },
+ }
+
+ if self.source_type == "EXTERNAL_SOURCES":
+ result = {
+ "type": self.source_type,
+ "externalSourcesConfiguration": {**base_config, "sources": self.sources},
+ }
+ return result
+
+ def execute(self, context: Context) -> Any:
+ self.validate_inputs()
+
+ result = self.hook.conn.retrieve_and_generate(
+ input={"text": self.input},
+ retrieveAndGenerateConfiguration=self.build_rag_config(),
+ **self.rag_kwargs,
+ )
+
+ self.log.info(
+ "\nPrompt: %s\nResponse: %s\nCitations: %s",
+ self.input,
+ result["output"]["text"],
+ result["citations"],
+ )
+ return result
+
+
+class BedrockRetrieveOperator(AwsBaseOperator[BedrockAgentRuntimeHook]):
+ """
+ Query a knowledge base and retrieve results with source citations.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:BedrockRetrieveOperator`
+
+ :param retrieval_query: The query to be made to the knowledge base. (templated)
+ :param knowledge_base_id: The unique identifier of the knowledge base that is queried. (templated)
+ Can only be specified if source_type='KNOWLEDGE_BASE'.
Review Comment:
I don't see source_type as a param for this operator?
##########
tests/system/providers/amazon/aws/example_bedrock_knowledge_base.py:
##########
@@ -480,6 +483,24 @@ def delete_opensearch_policies(collection_name: str):
)
# [END howto_sensor_bedrock_ingest_data]
+ # [START howto_operator_bedrock_retrieve_and_generate]
+ retrieve_and_generate = BedrockRaGOperator(
+ task_id="retrieve_and_generate",
+ input="Who was the CEO of Amazon on 2022?",
+ source_type="KNOWLEDGE_BASE",
Review Comment:
Should we have another one that tests the other source type?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Amazon Bedrock - Retrieve and RetrieveAndGenerate [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1598842673
##########
tests/system/providers/amazon/aws/example_bedrock_knowledge_base.py:
##########
@@ -480,6 +483,24 @@ def delete_opensearch_policies(collection_name: str):
)
# [END howto_sensor_bedrock_ingest_data]
+ # [START howto_operator_bedrock_retrieve_and_generate]
+ retrieve_and_generate = BedrockRaGOperator(
+ task_id="retrieve_and_generate",
+ input="Who was the CEO of Amazon on 2022?",
+ source_type="KNOWLEDGE_BASE",
Review Comment:
I couuuulllldd. The error is
```
ParamValidationError: Parameter validation failed:
Unknown parameter in retrieveAndGenerateConfiguration: "externalSourcesConfiguration", must be one of: type, knowledgeBaseConfiguration
```
I guess I could try/catch for that.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org