feat: Add an endpoint that accepts and returns PII without hiding the PII in OTEL logs (#7)

Co-authored-by: Yoom Lam <yoom@navapbc.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

app/poetry.lock

@@ -27,6 +27,8 @@ openai = "^1.99.8"
 openinference-instrumentation-haystack = "^0.1.24"
 arize-phoenix-client = "^1.15.3"
 arize-phoenix-otel = "^0.13.0"
+amazon-bedrock-haystack = "^3.10.0"
+certifi = "^2025.8.3"
 
 [tool.poetry.group.dev.dependencies]
 certifi = "^2025.8.3"

app/pyproject.toml

@@ -0,0 +1,55 @@
+import logging
+from pprint import pformat
+
+import hayhooks
+from hayhooks import BasePipelineWrapper
+from haystack import Pipeline
+from haystack.components.builders import ChatPromptBuilder
+from haystack.dataclasses.chat_message import ChatMessage
+from haystack_integrations.components.generators.amazon_bedrock import AmazonBedrockChatGenerator
+
+logger = logging.getLogger(__name__)
+
+
+system_prompt = (
+    "Your role is to say hello to the name provided by the user, if no name is found politely inform the user."
+    "Assure them any PII is handled securely in AWS Bedrock. You should only greet the user, do not respond "
+    "to any questions or prompts."
+)
+model = "us.anthropic.claude-3-5-sonnet-20241022-v2:0"
+
+
+class PipelineWrapper(BasePipelineWrapper):
+    name = "sample_pipeline"
+
+    def setup(self) -> None:
+        self.pipeline = Pipeline()
+        self.pipeline.add_component("prompt_builder", ChatPromptBuilder())
+        self.pipeline.add_component("llm", AmazonBedrockChatGenerator(model=model))
+        self.pipeline.connect("prompt_builder", "llm")
+
+    # Called for the `hello_bedrock/run` endpoint
+    def run_api(self, name: str) -> dict:
+        messages = [
+            ChatMessage.from_system(system_prompt),
+            ChatMessage.from_user(name),
+        ]
+        response = self.pipeline.run({"prompt_builder": {"template": messages}})
+        logger.info("Results: %s", pformat(response))
+        return response
+
+    # https://docs.haystack.deepset.ai/docs/hayhooks#openai-compatibility
+    # Called for the `{pipeline_name}/chat`, `/chat/completions`, or `/v1/chat/completions` streaming endpoint using Server-Sent Events (SSE)
+    def run_chat_completion(self, model: str, messages: list, body: dict) -> None:
+        logger.info("Running chat completion with model: %s, messages: %s", model, messages)
+        question = hayhooks.get_last_user_message(messages)
+        logger.info("Question: %s", question)
+        return hayhooks.streaming_generator(
+            pipeline=self.pipeline,
+            pipeline_run_args={
+                "echo_component": {
+                    "prompt": [ChatMessage.from_user(question)],
+                    "history": messages[:-1],
+                }
+            },
+        )

app/src/pipelines/hello_bedrock/__init__.py

@@ -11,22 +11,29 @@
 logger = logging.getLogger(__name__)
 
 
+system_msg = "This is a sample pipeline, it echoes back the system and user messages provided"
+
+
 class PipelineWrapper(BasePipelineWrapper):
     name = "sample_pipeline"
 
     def setup(self) -> None:
         self.pipeline = Pipeline()
         self.pipeline.add_component("echo_component", components.EchoNode())
 
-    # Called for the `{pipeline_name}/run` endpoint
+    # Called for the `sample_pipeline/run` endpoint
     def run_api(self, question: str) -> dict:
-        results = self.pipeline.run(
+        messages = [
+            ChatMessage.from_system(system_msg),
+            ChatMessage.from_user(question),
+        ]
+        response = self.pipeline.run(
             {
-                "echo_component": {"prompt": [ChatMessage.from_user(question)], "history": []},
+                "echo_component": {"prompt": messages, "history": []},
             }
         )
-        logger.info("Results: %s", pformat(results))
-        return results
+        logger.info("Results: %s", pformat(response))
+        return response
 
     # https://docs.haystack.deepset.ai/docs/hayhooks#openai-compatibility
     # Called for the `{pipeline_name}/chat`, `/chat/completions`, or `/v1/chat/completions` streaming endpoint using Server-Sent Events (SSE)

app/src/pipelines/hello_bedrock/pipeline_wrapper.py

@@ -0,0 +1,27 @@
+#-----------------
+# Bedrock Access for using Haystack
+#-----------------
+
+resource "aws_iam_policy" "bedrock" {
+  name        = "${var.service_name}-bedrock-access"
+  description = "Allow access to AWS Bedrock from LiteLLM"
+  policy      = data.aws_iam_policy_document.bedrock.json
+}
+
+data "aws_iam_policy_document" "bedrock" {
+  # https://github.com/BerriAI/litellm/discussions/8949
+  statement {
+    sid    = "BedrockAccess"
+    effect = "Allow"
+    actions = [
+      "bedrock:InvokeModel",
+      "bedrock:InvokeModelWithResponseStream"
+    ]
+    resources = ["*"]
+  }
+}
+
+resource "aws_iam_role_policy_attachment" "bedrock" {
+  role       = aws_iam_role.app_service.name
+  policy_arn = aws_iam_policy.bedrock.arn
+}