FEAT: add prompt suffix handling (#134)

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

Author: michelle-hadfield-nava

app/src/app_config.py

@@ -39,8 +39,9 @@ def db_session(self) -> db.Session:
     # so they are not unique across different Phoenix instances.
     PROMPT_VERSIONS: dict = {
         "extract_supports": "UHJvbXB0VmVyc2lvbjo0Ng==",
-        "generate_referrals": "UHJvbXB0VmVyc2lvbjo0OA==",
-        "generate_action_plan": "UHJvbXB0VmVyc2lvbjo1Mg==",
+        "generate_referrals": "UHJvbXB0VmVyc2lvbjo2MA==",  # if no suffix, the default Austin area prompt will be used
+        "generate_referrals_keystone": "UHJvbXB0VmVyc2lvbjo2Nw==",
+        "generate_action_plan": "UHJvbXB0VmVyc2lvbjo2OQ==",
         "crawl_gcta": "UHJvbXB0VmVyc2lvbjozNg==",
         "crawl_indeed": "UHJvbXB0VmVyc2lvbjozNA==",
     }

app/src/common/haystack_utils.py

@@ -6,8 +6,11 @@
 from src.common import phoenix_utils
 
 
-def get_phoenix_prompt(prompt_name: str, prompt_version_id: str = "") -> list[ChatMessage]:
-    prompt_ver = phoenix_utils.get_prompt_template(prompt_name, prompt_version_id)
+def get_phoenix_prompt(
+    prompt_name: str, prompt_version_id: str = "", suffix: str = ""
+) -> list[ChatMessage]:
+    full_prompt_name = f"{prompt_name}_{suffix}" if suffix else prompt_name
+    prompt_ver = phoenix_utils.get_prompt_template(full_prompt_name, prompt_version_id)
     return to_chat_messages(prompt_ver._template["messages"])
 
 

app/src/pipelines/generate_action_plan/pipeline_wrapper.py

@@ -48,11 +48,9 @@ def setup(self) -> None:
         pipeline = Pipeline()
         pipeline.add_component("llm", create_websearch())
 
-        prompt_template = haystack_utils.get_phoenix_prompt("generate_action_plan")
         pipeline.add_component(
             instance=ChatPromptBuilder(
-                template=prompt_template,
-                required_variables=["resources", "action_plan_json", "user_query"],
+                variables=["resources", "action_plan_json", "user_query"],
             ),
             name="prompt_builder",
         )
@@ -70,7 +68,10 @@ def setup(self) -> None:
 
     # Called for the `generate-action-plan/run` endpoint
     def run_api(
-        self, resources: list[Resource] | list[dict], user_email: str, user_query: str
+        self,
+        resources: list[Resource] | list[dict],
+        user_email: str,
+        user_query: str,
     ) -> dict:
         resource_objects = get_resources(resources)
 
@@ -87,14 +88,20 @@ def run_api(
                 return result
 
     def _run(self, resource_objects: list[Resource], user_email: str, user_query: str) -> dict:
+        prompt_template = haystack_utils.get_phoenix_prompt("generate_action_plan")
+
         response = self.pipeline.run(
             {
                 "logger": {
                     "messages_list": [
-                        {"resource_count": len(resource_objects), "user_email": user_email}
+                        {
+                            "resource_count": len(resource_objects),
+                            "user_email": user_email,
+                        }
                     ],
                 },
                 "prompt_builder": {
+                    "template": prompt_template,
                     "resources": format_resources(resource_objects),
                     "action_plan_json": action_plan_as_json,
                     "user_query": user_query,

app/src/pipelines/generate_referrals/pipeline_wrapper.py

@@ -100,14 +100,16 @@ def setup(self) -> None:
         self.pipeline = pipeline
 
     # Called for the `generate-referrals/run` endpoint
-    def run_api(self, query: str, user_email: str, prompt_version_id: str = "") -> dict:
+    def run_api(
+        self, query: str, user_email: str, prompt_version_id: str = "", suffix: str = ""
+    ) -> dict:
         with using_attributes(user_id=user_email), using_metadata({"user_id": user_email}):
             # Must set using_metadata context before calling tracer.start_as_current_span()
             assert isinstance(tracer, _tracers.OITracer), f"Got unexpected {type(tracer)}"
             with tracer.start_as_current_span(  # pylint: disable=not-context-manager,unexpected-keyword-arg
                 self.name, openinference_span_kind="chain"
             ) as span:
-                result = self._run(query, user_email, prompt_version_id)
+                result = self._run(query, user_email, prompt_version_id, suffix)
                 span.set_input(query)
                 try:
                     resp_obj = json.loads(result["llm"]["replies"][-1].text)
@@ -117,16 +119,18 @@ def run_api(self, query: str, user_email: str, prompt_version_id: str = "") -> d
                 span.set_status(Status(StatusCode.OK))
                 return result
 
-    def _run(self, query: str, user_email: str, prompt_version_id: str = "") -> dict:
-        # Retrieve the requested prompt_version_id and error if requested prompt version is not found
+    def _run(
+        self, query: str, user_email: str, prompt_version_id: str = "", suffix: str = ""
+    ) -> dict:
+        # Retrieve the requested prompt (with optional prompt_version_id and/or suffix)
         try:
             prompt_template = haystack_utils.get_phoenix_prompt(
-                "generate_referrals", prompt_version_id
+                "generate_referrals", prompt_version_id=prompt_version_id, suffix=suffix
             )
         except httpx.HTTPStatusError as he:
             raise HTTPException(
                 status_code=422,
-                detail=f"The requested prompt version '{prompt_version_id}' could not be retrieved due to HTTP status {he.response.status_code}",
+                detail=f"The requested prompt version '{prompt_version_id}' with suffix '{suffix}' could not be retrieved due to HTTP status {he.response.status_code}",
             ) from he
 
         try:

frontend/src/app/[locale]/generate-referrals/page.tsx

@@ -107,6 +107,7 @@ export default function Page() {
 
   async function findResources() {
     const prompt_version_id = searchParams?.get("prompt_version_id") ?? null;
+    const suffix = searchParams?.get("suffix") ?? undefined;
 
     setLoading(true);
     setResult(null);
@@ -118,6 +119,7 @@ export default function Page() {
         request,
         userEmail,
         prompt_version_id,
+        suffix,
       );
       setResourcesResultId(resultId);
       setErrorMessage(errorMessage);

frontend/src/components/ResourcesList.tsx

@@ -117,7 +117,9 @@ const ResourcesList = ({
           </button>
           <CardHeader className="p-3 ml-3">
             <CardTitle className="text-xl font-semibold text-gray-900 flex items-center gap-2">
-              <span className={`flex-shrink-0 w-7 h-7 ${getNumberBadgeClass(r.referral_type)} text-white rounded-full flex items-center justify-center text-m font-medium`}>
+              <span
+                className={`flex-shrink-0 w-7 h-7 ${getNumberBadgeClass(r.referral_type)} text-white rounded-full flex items-center justify-center text-m font-medium`}
+              >
                 {i + 1}
               </span>
               <div>{r.name}</div>

frontend/src/util/fetchActionPlan.ts

@@ -68,14 +68,20 @@ export async function fetchActionPlan(
   const timer = setTimeout(() => ac.abort(), 120_000);
 
   try {
+    const requestBody: {
+      resources: Resource[];
+      user_email: string;
+      user_query: string;
+    } = {
+      resources: resources,
+      user_email: userEmail,
+      user_query: userQuery,
+    };
+
     const upstream = await fetch(url, {
       method: "POST",
       headers,
-      body: JSON.stringify({
-        resources: resources,
-        user_email: userEmail,
-        user_query: userQuery,
-      }),
+      body: JSON.stringify(requestBody),
       cache: "no-store",
       signal: ac.signal,
     });

frontend/src/util/fetchResources.ts

@@ -5,6 +5,7 @@ export async function fetchResources(
   clientDescription: string,
   userEmail: string,
   prompt_version_id: string | null,
+  suffix?: string,
 ) {
   const apiDomain = await getApiDomain();
 
@@ -20,16 +21,26 @@ export async function fetchResources(
 
   const ac = new AbortController();
   const timer = setTimeout(() => ac.abort(), 600_000); // TODO make configurable
-  const requestBody = prompt_version_id
-    ? JSON.stringify({
-        query: clientDescription,
-        user_email: userEmail,
-        prompt_version_id: prompt_version_id,
-      })
-    : JSON.stringify({
-        query: clientDescription,
-        user_email: userEmail,
-      });
+
+  const requestData: {
+    query: string;
+    user_email: string;
+    prompt_version_id?: string;
+    suffix?: string;
+  } = {
+    query: clientDescription,
+    user_email: userEmail,
+  };
+
+  if (prompt_version_id) {
+    requestData.prompt_version_id = prompt_version_id;
+  }
+
+  if (suffix) {
+    requestData.suffix = suffix;
+  }
+
+  const requestBody = JSON.stringify(requestData);
 
   try {
     const upstream = await fetch(url, {