|
| 1 | +import argparse |
| 2 | +import functools |
| 3 | +import json |
| 4 | +import logging |
| 5 | +import os |
| 6 | +from pprint import pformat |
| 7 | +from typing import Any, Dict |
| 8 | + |
| 9 | +import requests |
| 10 | +from phoenix.client import Client |
| 11 | +from phoenix.client.experiments import run_experiment |
| 12 | +from phoenix.client.resources.datasets import Dataset |
| 13 | +from phoenix.client.resources.experiments.types import TaskOutput |
| 14 | + |
| 15 | +from src.common import phoenix_utils |
| 16 | +from src.pipelines.generate_referrals.pipeline_wrapper import PipelineWrapper |
| 17 | + |
| 18 | +logger = logging.getLogger(__name__) |
| 19 | + |
| 20 | + |
| 21 | +def get_sets(output: TaskOutput, expected: Dict[str, Any]) -> tuple[set[str], set[str]]: |
| 22 | + assert isinstance(output, list), f"Expected list of content, but got {type(output)}" |
| 23 | + assert len(output) == 1, f"Expected exactly one output, but got {len(output)}" |
| 24 | + output_obj = json.loads(output[0]["text"]) |
| 25 | + # Usually the output is {"resources": [...]}, sometimes it's just [...] |
| 26 | + resources = output_obj["resources"] if "resources" in output_obj else output_obj |
| 27 | + # Extract only the names of the resources from the output |
| 28 | + output_set = set([resource["name"] for resource in resources]) |
| 29 | + |
| 30 | + expected_json_str = expected[OUTPUT_COLUMN_NAME] |
| 31 | + expected_obj = json.loads(expected_json_str) |
| 32 | + expectation_set = set(expected_obj["expected_referrals"]) |
| 33 | + |
| 34 | + return output_set, expectation_set |
| 35 | + |
| 36 | + |
| 37 | +def recall(output: TaskOutput, expected: Dict[str, Any]) -> float: |
| 38 | + output_set, expectation_set = get_sets(output, expected) |
| 39 | + if len(expectation_set) == 0: |
| 40 | + raise ValueError("No expected referrals to compute recall.") |
| 41 | + return len(output_set.intersection(expectation_set)) / len(expectation_set) |
| 42 | + |
| 43 | + |
| 44 | +def precision(output: TaskOutput, expected: Dict[str, Any]) -> float: |
| 45 | + output_set, expectation_set = get_sets(output, expected) |
| 46 | + if len(output_set) == 0: |
| 47 | + return 0.0 |
| 48 | + return len(output_set.intersection(expectation_set)) / len(output_set) |
| 49 | + |
| 50 | + |
| 51 | +url_base = os.environ.get("DEPLOYED_API_URL") |
| 52 | +if url_base: |
| 53 | + logger.info("Using deployed API at %s", url_base) |
| 54 | + |
| 55 | + |
| 56 | +@functools.lru_cache |
| 57 | +def create_pipeline() -> PipelineWrapper: |
| 58 | + logger.info("Creating local Haystack pipeline") |
| 59 | + pipeline_wrapper = PipelineWrapper() |
| 60 | + pipeline_wrapper.setup() |
| 61 | + return pipeline_wrapper |
| 62 | + |
| 63 | + |
| 64 | +def get_question(example: dict) -> str: |
| 65 | + question_json_str = example["input"][INPUT_COLUMN_NAME] |
| 66 | + question_obj = json.loads(question_json_str) |
| 67 | + question = question_obj["caseworker_input"] |
| 68 | + return question |
| 69 | + |
| 70 | + |
| 71 | +def query_pipeline(example: dict) -> TaskOutput: |
| 72 | + question = get_question(example) |
| 73 | + logger.info("Getting answer for: %r", question) |
| 74 | + response = create_pipeline().run_api(query=question) |
| 75 | + replies = response["llm"]["replies"] |
| 76 | + assert len(replies) == 1, f"Expected exactly one reply but got {len(replies)}" |
| 77 | + return replies[0].to_dict()["content"] |
| 78 | + |
| 79 | + |
| 80 | +def query_api(example: dict) -> TaskOutput: |
| 81 | + question = get_question(example) |
| 82 | + logger.info("Getting answer for: %r", question) |
| 83 | + |
| 84 | + assert url_base, "DEPLOYED_API_URL is not set -- add it to override.env" |
| 85 | + response = requests.post( |
| 86 | + f"{url_base}/generate_referrals/run", |
| 87 | + headers={ |
| 88 | + "accept": "application/json", |
| 89 | + "Content-Type": "application/json", |
| 90 | + }, |
| 91 | + json={"query": question}, |
| 92 | + timeout=60, |
| 93 | + ) |
| 94 | + resp_obj = response.json() |
| 95 | + logger.info("Response: %s", pformat(resp_obj, width=160)) |
| 96 | + assert len(resp_obj["result"]["llm"]["replies"]) == 1, "Expected exactly one reply" |
| 97 | + return resp_obj["result"]["llm"]["replies"][0]["_content"] |
| 98 | + |
| 99 | + |
| 100 | +def run(dataset: Dataset, client: Client) -> None: |
| 101 | + task = query_api if url_base else query_pipeline |
| 102 | + logger.info("Running experiment using %r", task.__name__) |
| 103 | + run_experiment( |
| 104 | + dataset=dataset, |
| 105 | + task=task, |
| 106 | + evaluators=[recall, precision], |
| 107 | + client=client, |
| 108 | + ) |
| 109 | + |
| 110 | + |
| 111 | +INPUT_COLUMN_NAME = "Input" |
| 112 | +OUTPUT_COLUMN_NAME = "Output" |
| 113 | + |
| 114 | + |
| 115 | +def export_dataset(dataset_name: str, filename: str, client: Client) -> None: |
| 116 | + """Retrieve a dataset by name and save it as a CSV file. Useful for exporting a dataset from the deployed Phoenix.""" |
| 117 | + try: |
| 118 | + dataset = client.datasets.get_dataset(dataset=dataset_name) |
| 119 | + # Convert the dataset to a pandas DataFrame to save as CSV |
| 120 | + df = dataset.to_dataframe() |
| 121 | + |
| 122 | + # Remove the extraneous 'input' and 'output' keys and just keep the values |
| 123 | + # This ensures that when the CSV file is imported, the input and output columns are the same as in the deployed Phoenix |
| 124 | + df[INPUT_COLUMN_NAME] = df["input"].apply(lambda x: x[INPUT_COLUMN_NAME]) |
| 125 | + df[OUTPUT_COLUMN_NAME] = df["output"].apply(lambda x: x[OUTPUT_COLUMN_NAME]) |
| 126 | + # Convert metadata to JSON string so it can be correctly parsed when running experiments if needed |
| 127 | + df["metadata"] = df["metadata"].apply(lambda x: json.dumps(x)) |
| 128 | + df = df[[INPUT_COLUMN_NAME, OUTPUT_COLUMN_NAME, "metadata"]] |
| 129 | + df.to_csv(filename, index=False) |
| 130 | + except ValueError as e: |
| 131 | + logger.error("Error retrieving dataset: %s", e) |
| 132 | + print_datasets(client) |
| 133 | + |
| 134 | + |
| 135 | +def print_datasets(client: Client) -> list: |
| 136 | + all_datasets = client.datasets.list() |
| 137 | + logger.info("%d available datasets:\n%s", len(all_datasets), pformat(all_datasets)) |
| 138 | + return all_datasets |
| 139 | + |
| 140 | + |
| 141 | +def import_dataset(client: Client, filename: str, dataset_name: str) -> Any: |
| 142 | + """ |
| 143 | + Import a dataset from a CSV file into Phoenix. |
| 144 | + Useful for loading dataset locally to test experiments. |
| 145 | + Reminder to reduce the number of rows in the CSV file for development. |
| 146 | + """ |
| 147 | + return client.datasets.create_dataset( |
| 148 | + name=dataset_name, |
| 149 | + csv_file_path=filename, |
| 150 | + input_keys=[INPUT_COLUMN_NAME], |
| 151 | + output_keys=[OUTPUT_COLUMN_NAME], |
| 152 | + # Caution: Note that the metadata column is different than on the deployed Phoenix |
| 153 | + # In the Phoenix UI, click on each example to compare -- don't compare in the table view, |
| 154 | + # which shows multiple examples and hides the "Input" and "Output" keys |
| 155 | + metadata_keys=["metadata"], |
| 156 | + ) |
| 157 | + |
| 158 | + |
| 159 | +def main() -> None: |
| 160 | + logging.basicConfig(format="%(levelname)s - %(name)s - %(message)s", level=logging.INFO) |
| 161 | + |
| 162 | + parser = argparse.ArgumentParser() |
| 163 | + parser.add_argument("dataset", type=str, default="brandon2") |
| 164 | + parser.add_argument( |
| 165 | + "action", |
| 166 | + type=str, |
| 167 | + choices=["export", "import", "run", "run_on_deployed"], |
| 168 | + default="run", |
| 169 | + ) |
| 170 | + args = parser.parse_args() |
| 171 | + |
| 172 | + logger.info("Action=%s Dataset=%r", args.action, args.dataset) |
| 173 | + |
| 174 | + if args.action == "export": |
| 175 | + client_to_deployed_phx = phoenix_utils.client_to_deployed_phoenix() |
| 176 | + export_dataset(args.dataset, "dataset.csv", client_to_deployed_phx) |
| 177 | + logger.info("Export complete. Check dataset.csv file.") |
| 178 | + return |
| 179 | + |
| 180 | + client = phoenix_utils._create_client() |
| 181 | + if args.action == "import": |
| 182 | + import_dataset(client, "dataset.csv", args.dataset) |
| 183 | + logger.info("Import complete. Check Phoenix UI.") |
| 184 | + return |
| 185 | + |
| 186 | + if args.action == "run_on_deployed": |
| 187 | + logger.info("Running experiment on dataset in deployed Phoenix") |
| 188 | + client = phoenix_utils.client_to_deployed_phoenix() |
| 189 | + |
| 190 | + try: |
| 191 | + # Get the dataset, run the experiment, and post the results |
| 192 | + dataset = client.datasets.get_dataset(dataset=args.dataset) |
| 193 | + run(dataset, client) |
| 194 | + logger.info("Check results in the Phoenix UI.") |
| 195 | + except ValueError as e: |
| 196 | + logger.error("Error retrieving dataset: %s", e) |
| 197 | + print_datasets(client) |
0 commit comments