feat: Lookup prompt template from Arize Phoenix (#23)

Author: yoomlam

app/Makefile

@@ -227,3 +227,6 @@ help: ## Prints the help documentation and info about each command
 # Usage: make extract-supports NAME="SupportName" FILE="path/to/file"
 extract-supports:
 	$(PY_RUN_CMD) extract-supports "$(NAME)" "$(FILE)"
+
+copy-prompts:
+	$(PY_RUN_CMD) copy-prompts

app/README.md

@@ -81,3 +81,11 @@ Based on [documentation](https://arize.com/docs/phoenix/self-hosting/features/au
 ```
       - PHOENIX_API_KEY=<paste API key>
 ```
+
+### Copying prompts from the deployed Phoenix
+
+For local development, you'll likely want to replicate the prompts in the deployed Phoenix instance onto your local Phoenix instance.
+To do so, add the `DEPLOYED_PHOENIX_URL` and `DEPLOYED_PHOENIX_API_KEY` environment variables to `override.env`.
+Create a system API key at `$DEPLOYED_PHOENIX_URL/settings/general`.
+Then run `make copy-prompts`. This will copy the prompt versions specified in `app_config.py`, which are the ones used in the deployed app.
+Remember to do this every time the prompt version is updated in `app_config.py`. When running locally, the latest version of the prompt is used.

app/docker-compose.yml

@@ -38,8 +38,6 @@ services:
       - PHOENIX_TLS_KEY_FILE_PASSWORD=
       # For verifying client certificates
       - PHOENIX_TLS_VERIFY_CLIENT=False
-      # In case we want to disable PII Redaction
-      - REDACT_PII=${REDACT_PII:-True}
 
       # See README.md to enable authentication locally
       # - PHOENIX_ENABLE_AUTH=True

app/local.env

@@ -100,3 +100,6 @@ BATCH_OTEL=false
 # local Phoenix instance within the Docker compose network
 PHOENIX_COLLECTOR_ENDPOINT=https://phoenix:6006
 PHOENIX_PROJECT_NAME=local-docker-project
+
+# Disable PII redaction for local development
+# REDACT_PII=False

app/pyproject.toml

@@ -63,6 +63,7 @@ db-migrate = "src.db.migrations.run:up"
 db-migrate-down = "src.db.migrations.run:down"
 db-migrate-down-all = "src.db.migrations.run:downall"
 extract-supports= "src.ingestion.extract_supports:main"
+copy-prompts= "src.common.phoenix_utils:copy_deployed_prompts"
 
 [tool.black]
 line-length = 100

app/src/app_config.py

@@ -5,6 +5,7 @@
 
 
 class AppConfig(PydanticBaseEnvConfig):
+    environment: str = "local"
     # Set HOST to 127.0.0.1 by default to avoid other machines on the network
     # from accessing the application. This is especially important if you are
     # running the application locally on a public network. This needs to be
@@ -24,5 +25,14 @@ def db_client(self) -> db.PostgresDBClient:
     def db_session(self) -> db.Session:
         return self.db_client.get_session()
 
+    # These versions should only be used for the deployed Phoenix instance.
+    # Version ids are base64 encodings of 'PromptVersion:N' where N is simply a counter,
+    # so they are not unique across different Phoenix instances.
+    PROMPT_VERSIONS: dict = {
+        "sample_rag": "UHJvbXB0VmVyc2lvbjox",
+        "extract_supports": "UHJvbXB0VmVyc2lvbjoz",
+        "generate_referrals": "UHJvbXB0VmVyc2lvbjo0",
+    }
+
 
 config = AppConfig()

app/src/common/haystack_utils.py

@@ -0,0 +1,48 @@
+from typing import Sequence
+
+from haystack.dataclasses.chat_message import ChatMessage
+from phoenix.client.__generated__ import v1
+
+from src.common import phoenix_utils
+
+
+def get_phoenix_prompt(prompt_name: str) -> list[ChatMessage]:
+    prompt_ver = phoenix_utils.get_prompt_template(prompt_name)
+    return to_chat_messages(prompt_ver._template["messages"])
+
+
+def to_chat_messages(
+    msg_list: Sequence[dict | v1.PromptMessage | ChatMessage],
+) -> list[ChatMessage]:
+    """Convert a list of dicts or Phoenix PromptMessage to a list of Haystack ChatMessage."""
+    messages = []
+    for msg in msg_list:
+        if isinstance(msg, ChatMessage):
+            messages.append(msg)
+            continue
+        elif not isinstance(msg, dict):  # PromptMessage is a TypedDict
+            raise ValueError(f"Expected dict or ChatMessage, got {type(msg)}")
+
+        role = msg["role"]
+        content = msg["content"]
+
+        assert isinstance(content, list), f"Expected list content, got {type(content)}: {content}"
+        assert len(content) == 1, f"Expected single content, got {len(content)} items: {content}"
+        assert content[0]["type"] == "text", f"Expected text content, got {content[0]['type']}"
+        assert "text" in content[0], f"Expected 'text' in content[0], got {content[0]}"
+        text = content[0]["text"]
+
+        if role == "system":
+            assert isinstance(text, str), f"Expected string, got {type(text)}"
+            chat_msg = ChatMessage.from_system(text)
+        elif role == "user":
+            assert isinstance(text, str), f"Expected string, got {type(text)}"
+            chat_msg = ChatMessage.from_user(text)
+        elif role == "assistant":
+            assert isinstance(text, str), f"Expected string, got {type(text)}"
+            chat_msg = ChatMessage.from_assistant(text)
+        else:
+            raise ValueError(f"Unexpected role: {role} for message {msg}")
+        messages.append(chat_msg)
+
+    return messages

app/src/common/phoenix_utils.py

@@ -1,26 +1,28 @@
 import logging
 import os
+from pprint import pformat
 
 import httpx
 import opentelemetry.exporter.otlp.proto.http.trace_exporter as otel_trace_exporter
+import phoenix.otel
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
 
 # https://docs.arize.com/phoenix/tracing/integrations-tracing/haystack
 # Arize's Phoenix observability platform
-import phoenix.client
-import phoenix.otel
-from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from phoenix.client import Client
+from phoenix.client.types import PromptVersion
 
 from src.app_config import config
 from src.logging.presidio_pii_filter import PresidioRedactionSpanProcessor
 
 logger = logging.getLogger(__name__)
 
 
-def _create_client() -> phoenix.client.Client:
-    logger.info("Creating Phoenix client to %s", config.phoenix_collector_endpoint)
-    # If base_url is None, then phoenix.client.Client defaults to PHOENIX_COLLECTOR_ENDPOINT
-    # env variable value or "http://localhost:6006"
-    return phoenix.client.Client(base_url=config.phoenix_collector_endpoint)
+def _create_client(
+    url: str = config.phoenix_collector_endpoint, api_key: str | None = None
+) -> Client:
+    logger.info("Creating Phoenix client to %s", url)
+    return Client(base_url=url, api_key=api_key)
 
 
 def service_alive() -> bool:
@@ -70,3 +72,73 @@ def configure_phoenix(only_if_alive: bool = True) -> None:
         if config.batch_otel:
             tracer_provider.add_span_processor(BatchSpanProcessor(span_exporter))
         tracer_provider.add_span_processor(pii_processor)
+
+
+def get_prompt_template(prompt_name: str) -> PromptVersion:
+    """Retrieve a prompt template from Phoenix by name.
+    https://arize.com/docs/phoenix/sdk-api-reference/python/overview#prompt-management
+    """
+    prompt_params = which_prompt_version(prompt_name)
+    client = _create_client()
+    prompt = client.prompts.get(**prompt_params)
+    logger.info(
+        "Retrieved prompt with %r: id='%s'\n%s", prompt_params, prompt.id, pformat(prompt._dumps())
+    )
+    return prompt
+
+
+def which_prompt_version(prompt_name: str) -> dict:
+    if config.environment == "local":
+        # Get the latest version regardless of tags
+        return {"prompt_identifier": prompt_name}
+
+    # Use the hardcoded version ids
+    return {"prompt_version_id": config.PROMPT_VERSIONS[prompt_name]}
+
+
+def copy_deployed_prompts() -> None:
+    logging.basicConfig(format="%(levelname)s - %(name)s -  %(message)s", level=logging.INFO)
+
+    url = os.environ.get("DEPLOYED_PHOENIX_URL")
+    api_key = os.environ.get("DEPLOYED_PHOENIX_API_KEY")
+    logger.info("Copying prompts from %s with API key: %r", url, api_key)
+    assert url, "DEPLOYED_PHOENIX_URL is not set -- add it to override.env"
+    assert api_key, "DEPLOYED_PHOENIX_API_KEY is not set -- add it to override.env"
+
+    src_client = _create_client(url, api_key=api_key)
+    local_client = _create_client()
+    for prompt in list_prompts(src_client):
+        # The prompt id is base64 encoding of 'Prompt:N' where N is simply a counter
+        logger.info("Copying prompt: %r with id=%r)", prompt["name"], prompt["id"])
+        copy_prompt(src_client, local_client, prompt["name"])
+
+
+def list_prompts(client: Client) -> list[dict]:
+    "client.prompts doesn't have a list() method, so use the underlying httpx client."
+    response = client._client.get("/v1/prompts")
+    return response.json()["data"]
+
+
+def copy_prompt(src_client: Client, local_client: Client, prompt_name: str) -> None:
+    "Copy a prompt from src_client to local_client"
+    if prompt_name not in config.PROMPT_VERSIONS:
+        logger.warning("No version id found for prompt %r -- skipping", prompt_name)
+        return
+
+    prompt_ver = src_client.prompts.get(prompt_version_id=config.PROMPT_VERSIONS[prompt_name])
+    logger.info("Retrieved prompt with id='%s'\n%s", prompt_ver.id, pformat(prompt_ver._dumps()))
+
+    logger.info("Creating prompt %r in %r", prompt_name, local_client._client.base_url)
+    # If prompt_name already exists, a new prompt version will be created
+    local_client.prompts.create(
+        version=prompt_ver, name=prompt_name, prompt_description=prompt_ver._description
+    )
+
+
+def list_prompt_version_ids(prompt_name: str, client: Client) -> list[str]:
+    "List all version ids for a given prompt name. client.prompts doesn't have a list_versions() method."
+    response = client._client.get(f"/v1/prompts/{prompt_name}/versions")
+    resp_data = response.json()["data"]
+    # version tags are not in the response
+    return [ver["id"] for ver in resp_data]
+    # To get tags for the version: client.prompts.tags.list(prompt_version_id=version_id)

app/src/ingestion/extract_supports.py

@@ -15,13 +15,13 @@
 from haystack.components.converters import PyPDFToDocument
 from haystack.components.preprocessors import DocumentSplitter
 from haystack.dataclasses import Document
-from haystack.dataclasses.chat_message import ChatMessage
 from haystack_integrations.components.generators.amazon_bedrock import AmazonBedrockChatGenerator
 from pydantic import BaseModel, Field
 from smart_open import open as smart_open
 
 from src.adapters import db
 from src.app_config import config
+from src.common import haystack_utils
 from src.db.models.support_listing import Support, SupportListing
 
 logger = logging.getLogger(__name__)
@@ -76,21 +76,7 @@ class SupportEntry(BaseModel):
     description: str | None = Field(description="2-sentence summary, including offerings")
 
 
-SYSTEM_PROMPT_TEMPLATE = f"""Using only the document content provided by the user, return a JSON list of objects.
-Each object must match exactly this schema:
-```
-{json.dumps(SupportEntry.model_json_schema(), indent=2)}
-```
-
-Rules:
-- Output ONLY raw JSON (no markdown fences, no commentary).
-- If a field is missing in the PDF, use null or [] as appropriate.
-- Keep strings concise; avoid line breaks inside values.
-"""
-
-USER_TEMPLATE = """Document content:
-{{ doc.content }}
-"""
+OUTPUT_SCHEMA = json.dumps(SupportEntry.model_json_schema(), indent=2)
 
 
 def create_llm() -> AmazonBedrockChatGenerator:  # pragma: no cover
@@ -105,12 +91,9 @@ def create_llm() -> AmazonBedrockChatGenerator:  # pragma: no cover
 def build_pipeline() -> AsyncPipeline:
     pipe = AsyncPipeline()
 
-    messages = [
-        ChatMessage.from_system(SYSTEM_PROMPT_TEMPLATE),
-        ChatMessage.from_user(USER_TEMPLATE),
-    ]
+    chat_template = haystack_utils.get_phoenix_prompt("extract_supports")
     pipe.add_component(
-        "prompt_builder", ChatPromptBuilder(template=messages, required_variables="*")
+        "prompt_builder", ChatPromptBuilder(template=chat_template, required_variables="*")
     )
     pipe.add_component("llm", create_llm())
     # If needed, add OutputValidator to retry the LLM call -- https://haystack.deepset.ai/tutorials/28_structured_output_with_loop
@@ -124,7 +107,7 @@ async def run_pipeline(pipeline: AsyncPipeline, doc: Document) -> list[dict]:
     assert doc.content
     logger.info("Running pipeline with subdoc content length: %d", len(doc.content))
 
-    _result = await pipeline.run_async({"prompt_builder": {"doc": doc}})
+    _result = await pipeline.run_async({"prompt_builder": {"schema": OUTPUT_SCHEMA, "doc": doc}})
     assert len(_result["llm"]["replies"]) == 1
     reply = _result["llm"]["replies"][0]
     # Useful info for checking if tokens have reached the limit for the LLM
@@ -186,8 +169,8 @@ def save_to_db(
 
     support_listing_id = existing_listing.id if existing_listing else support_listing.id
     assert support_listing_id
-    # Populate support records
     for support in support_entries:
+        logger.info("Adding Support record: %r", support.name)
         support_record = Support(
             support_listing_id=support_listing_id,
             name=support.name,
@@ -216,9 +199,6 @@ def main() -> None:  # pragma: no cover
     logger.info("Extracted content length: %d", len(doc.content))
     extracted_supports = extract_support_entries(args.name, doc)
 
-    for support in extracted_supports.values():
-        logger.info("Support: %r", support.name)
-
     with config.db_session() as db_session, db_session.begin():
         support_listing = SupportListing(name=args.name, uri=args.filepath)
         save_to_db(db_session, support_listing, extracted_supports.values())

app/src/pipelines/generate_referrals/pipeline_wrapper.py

@@ -12,6 +12,7 @@
 from sqlalchemy.inspection import inspect
 
 from src.app_config import config
+from src.common import haystack_utils
 from src.db.models.support_listing import Support
 
 logger = logging.getLogger(__name__)
@@ -26,46 +27,17 @@ class Resource(BaseModel):
 
 
 resource_as_json = json.dumps(Resource.model_json_schema(), indent=2)
-
-system_prompt = """
-You are a supporting API for Goodwill Central Texas Referral. You are designed to help career case managers provide high-quality, local resource referrals to client's in Central Texas.
-Your role is to support Goodwill Central Texas career case managers working with low-income job seekers and learners in Austin and surrounding counties (Bastrop, Blanco, Burnet, Caldwell, DeWitt, Fayette, Gillespie, Gonzales, Hays, Lavaca, Lee, Llano, Mason, Travis, Williamson).
-
-## Task Checklist
-    - Evaluate the client needs and determine their eligibility (Factors to consider: age, income, disability, immigration/veteran status, number of dependents)
-    - Prioritize Goodwill resources first (Basic Needs Resource packet, Goodwill websites)
-    - Rank recommendations by proximity, eligibility fit, and other relevant factors
-
-## Core Instructions
-    - Use only trusted and up-to-date sources: Goodwill, government, vetted nonprofits, trusted news outlets (Findhelp, 211, Connect ATX permitted). Never use unreliable websites (e.g., shelterlistings.org, needhelppayingbills.com).
-    - Never invent or fabricate resources. If none are available, state this clearly and suggest actionable, specific next steps
-
-List of resources to choose from:
-{% for s in supports %}
-   - {{ s.content }}
-{% endfor %}
-
-## Response Constraints
-    - Your response should ONLY include resources.
-    - Do not summarize your assessment of the clients needs.
-    - Limit the description for a resource to be less than 255 words.
-    - Return a JSON list of resources in the following format:
-        '''{{ resource_json }}'''
-    """
 model = "us.anthropic.claude-3-5-sonnet-20241022-v2:0"
 
-prompt_template = [
-    ChatMessage.from_system(system_prompt),
-    ChatMessage.from_user("""User query: {{ query }}"""),
-]
-
 
 class PipelineWrapper(BasePipelineWrapper):
     name = "generate_referrals"
 
     def setup(self) -> None:
         pipeline = Pipeline()
         pipeline.add_component("llm", AmazonBedrockChatGenerator(model=model))
+
+        prompt_template = haystack_utils.get_phoenix_prompt("generate_referrals")
         pipeline.add_component(
             instance=ChatPromptBuilder(
                 template=prompt_template, required_variables=["query", "supports", "resource_json"]