ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter2/prompt-engineering/tau_bench/model_utils/model/anyscale.py

@@ -0,0 +1,89 @@
+import os
+
+from tau_bench.model_utils.api.datapoint import Datapoint
+from tau_bench.model_utils.model.chat import ChatModel, Message
+from tau_bench.model_utils.model.completion import approx_cost_for_datapoint, approx_prompt_str
+from tau_bench.model_utils.model.general_model import wrap_temperature
+from tau_bench.model_utils.model.utils import approx_num_tokens
+
+API_KEY_ENV_VAR = "ANYSCALE_API_KEY"
+BASE_URL = "https://api.endpoints.anyscale.com/v1"
+
+PRICE_PER_INPUT_TOKEN_MAP = {"meta-llama/Meta-Llama-3-8B-Instruct": ...}
+INPUT_PRICE_PER_TOKEN_FALLBACK = 10 / 1000000
+
+CAPABILITY_SCORE_MAP = {
+    "meta-llama/Meta-Llama-3-8B-Instruct": 0.2,
+    "meta-llama/Meta-Llama-3-70B-Instruct": 0.6,
+}
+CAPABILITY_SCORE_FALLBACK = 0.2
+
+# TODO: implement
+LATENCY_MS_PER_OUTPUT_TOKEN_MAP = {}
+# TODO: implement
+LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK = 0.0
+
+MAX_CONTEXT_LENGTH_MAP = {
+    "meta-llama/Meta-Llama-3-8B-Instruct": 8192,
+    "meta-llama/Meta-Llama-3-70B-Instruct": 8192,
+}
+MAX_CONTEXT_LENGTH_FALLBACK = 8192
+
+
+class AnyscaleModel(ChatModel):
+    def __init__(
+        self,
+        model: str,
+        api_key: str | None = None,
+        temperature: float = 0.0,
+    ) -> None:
+        from openai import AsyncOpenAI, OpenAI
+
+        self.model = model
+
+        api_key = None
+        if api_key is None:
+            api_key = os.getenv(API_KEY_ENV_VAR)
+            if api_key is None:
+                raise ValueError(f"{API_KEY_ENV_VAR} environment variable is not set")
+        self.client = OpenAI(api_key=api_key, base_url=BASE_URL)
+        self.async_client = AsyncOpenAI(api_key=api_key, base_url=BASE_URL)
+        self.temperature = temperature
+
+    def generate_message(
+        self,
+        messages: list[Message],
+        force_json: bool,
+        temperature: float | None = None,
+    ) -> Message:
+        if temperature is None:
+            temperature = self.temperature
+        msgs = self.build_generate_message_state(messages)
+        res = self.client.chat.completions.create(
+            model=self.model,
+            messages=msgs,
+            temperature=wrap_temperature(temperature),
+            response_format={"type": "json_object" if force_json else "text"},
+        )
+        return self.handle_generate_message_response(
+            prompt=msgs, content=res.choices[0].message.content, force_json=force_json
+        )
+
+    def get_approx_cost(self, dp: Datapoint) -> float:
+        cost_per_token = PRICE_PER_INPUT_TOKEN_MAP.get(self.model, INPUT_PRICE_PER_TOKEN_FALLBACK)
+        return approx_cost_for_datapoint(dp=dp, price_per_input_token=cost_per_token)
+
+    def get_latency(self, dp: Datapoint) -> float:
+        latency_per_output_token = LATENCY_MS_PER_OUTPUT_TOKEN_MAP.get(
+            self.model, LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK
+        )
+        return approx_cost_for_datapoint(dp=dp, price_per_input_token=latency_per_output_token)
+
+    def get_capability(self) -> float:
+        return CAPABILITY_SCORE_MAP.get(self.model, CAPABILITY_SCORE_FALLBACK)
+
+    def supports_dp(self, dp: Datapoint) -> bool:
+        prompt = approx_prompt_str(dp)
+        return approx_num_tokens(prompt) <= MAX_CONTEXT_LENGTH_MAP.get(
+            self.model, MAX_CONTEXT_LENGTH_FALLBACK
+        )

chapter2/prompt-engineering/tau_bench/model_utils/model/chat.py

@@ -0,0 +1,608 @@
+import abc
+import enum
+import json
+from typing import Any, TypeVar
+
+from pydantic import BaseModel
+
+from tau_bench.model_utils.api.datapoint import (
+    BinaryClassifyDatapoint,
+    ClassifyDatapoint,
+    Datapoint,
+    GenerateDatapoint,
+    ParseDatapoint,
+    ParseForceDatapoint,
+    ScoreDatapoint,
+)
+from tau_bench.model_utils.api.types import PartialObj
+from tau_bench.model_utils.model.exception import ModelError
+from tau_bench.model_utils.model.general_model import GeneralModel
+from tau_bench.model_utils.model.utils import (
+    add_md_tag,
+    clean_top_level_keys,
+    display_choices,
+    json_response_to_obj_or_partial_obj,
+    optionalize_type,
+    parse_json_or_json_markdown,
+    try_classify_recover,
+    type_to_json_schema_string,
+)
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class Role(str, enum.Enum):
+    SYSTEM = "system"
+    ASSISTANT = "assistant"
+    USER = "user"
+
+
+class Message(BaseModel):
+    role: Role
+    content: str
+    obj: dict[str, Any] | None = None
+
+    def model_dump(self, **kwargs) -> dict[str, Any]:
+        if self.obj is not None:
+            return super().model_dump(**kwargs)
+        return {"role": self.role, "content": self.content}
+
+
+class PromptSuffixStrategy(str, enum.Enum):
+    JSON = "json"
+    JSON_MD_BLOCK = "json_md_block"
+
+
+def force_json_prompt(
+    text: str,
+    suffix_strategy: PromptSuffixStrategy = PromptSuffixStrategy.JSON,
+) -> str:
+    if suffix_strategy == PromptSuffixStrategy.JSON:
+        return f"{text}\n\nValid JSON:"
+    elif suffix_strategy == PromptSuffixStrategy.JSON_MD_BLOCK:
+        return f'{text}\n\nThe result should be a valid JSON object (according to the definition in the provided schema) in a markdown block only. For example:\nassistant:```json\n{{"items": ["value"]}}\n```'
+    else:
+        raise ValueError(f"Invalid suffix strategy: {suffix_strategy}")
+
+
+def build_generate_state(
+    instruction: str,
+    text: str,
+    examples: list[GenerateDatapoint] | None = None,
+) -> list[Message]:
+    messages = []
+    if examples is not None:
+        for example in examples:
+            example_msgs = [
+                Message(role=Role.SYSTEM, content=example.instruction),
+                Message(role=Role.USER, content=example.text),
+                Message(role=Role.ASSISTANT, content=example.response),
+            ]
+            messages.extend(example_msgs)
+    messages.append(Message(role=Role.SYSTEM, content=instruction))
+    messages.append(Message(role=Role.USER, content=text))
+    return messages
+
+
+def build_parse_force_state(
+    instruction: str,
+    typ: type[T] | dict[str, Any],
+    text: str | None = None,
+    examples: list[ParseForceDatapoint] | None = None,
+    suffix_strategy: PromptSuffixStrategy = PromptSuffixStrategy.JSON,
+) -> list[Message]:
+    def display_sample(
+        instr: str,
+        ty: type[T] | dict[str, Any],
+        t: str | None = None,
+        response: T | dict[str, Any] | None = None,
+    ) -> Message | list[Message]:
+        if isinstance(ty, dict):
+            json_schema_string = json.dumps(ty)
+        else:
+            json_schema_string = type_to_json_schema_string(ty)
+        text_insert = "" if t is None else f"\n\nText:\n{t}"
+        input_text = force_json_prompt(
+            text=f"Instruction:\n{instr}{text_insert}\n\nSchema:\n{json_schema_string}",
+            suffix_strategy=suffix_strategy,
+        )
+        if response is not None:
+            if isinstance(response, dict):
+                response_display = json.dumps(response)
+            else:
+                response_display = json.dumps(response.model_dump())
+            return [
+                Message(role=Role.USER, content=input_text),
+                Message(role=Role.ASSISTANT, content=response_display),
+            ]
+        else:
+            return Message(role=Role.USER, content=input_text)
+
+    messages = [
+        Message(
+            role=Role.SYSTEM,
+            content="Generate an object with the provided instruction, text, and schema.",
+        ),
+    ]
+    if examples is not None:
+        for example in examples:
+            example_msgs = display_sample(
+                instr=example.instruction,
+                ty=example.typ,
+                t=example.text,
+                response=example.response,
+            )
+            assert isinstance(example_msgs, list) and all(
+                isinstance(msg, Message) for msg in example_msgs
+            )
+            messages.extend(example_msgs)
+    messages.append(display_sample(instr=instruction, ty=typ, t=text))
+    return messages
+
+
+def build_score_state(
+    instruction: str,
+    text: str,
+    min: int,
+    max: int,
+    examples: list[ScoreDatapoint] | None = None,
+    suffix_strategy: PromptSuffixStrategy = PromptSuffixStrategy.JSON,
+) -> list[Message]:
+    def display_sample(
+        instr: str, t: str, mn: int, mx: int, response: int | None = None
+    ) -> list[Message] | Message:
+        if mn > mx:
+            raise ValueError(f"Invalid range: [{mn}, {mx}]")
+        input_text = force_json_prompt(
+            f"Instruction:\n{instr}\n\nText:\n{t}\n\nRange:\n[{mn}, {mx}]",
+            suffix_strategy,
+        )
+        if response is not None:
+            return [
+                Message(role=Role.USER, content=input_text),
+                Message(role=Role.ASSISTANT, content=f'{{"score": {response}}}'),
+            ]
+        else:
+            return Message(role=Role.USER, content=input_text)
+
+    messages = [
+        Message(
+            role=Role.SYSTEM,
+            content='Score the following text with the provided instruction and range as an integer value in valid JSON:\n{"score": number}',
+        ),
+    ]
+    if examples is not None:
+        for example in examples:
+            example_msgs = display_sample(
+                instr=example.instruction,
+                t=example.text,
+                mn=example.min,
+                mx=example.max,
+                response=example.response,
+            )
+            assert isinstance(example_msgs, list) and all(
+                isinstance(msg, Message) for msg in example_msgs
+            ), example_msgs
+            messages.extend(example_msgs)
+    messages.append(display_sample(instr=instruction, t=text, mn=min, mx=max))
+    return messages
+
+
+def build_parse_state(
+    text: str,
+    typ: type[T] | dict[str, Any],
+    examples: list[ParseDatapoint] | None = None,
+    suffix_strategy: PromptSuffixStrategy = PromptSuffixStrategy.JSON,
+) -> list[Message]:
+    def display_sample(
+        t: str,
+        ty: type[T] | dict[str, Any],
+        response: T | PartialObj | dict[str, Any] | None = None,
+    ) -> Message | list[Message]:
+        if isinstance(ty, dict):
+            json_schema_string = json.dumps(ty)
+        else:
+            optionalized_typ = optionalize_type(ty)
+            json_schema_string = type_to_json_schema_string(optionalized_typ)
+        input_text = force_json_prompt(
+            f"Text:\n{t}\n\nSchema:\n{json_schema_string}",
+            suffix_strategy=suffix_strategy,
+        )
+        if response is not None:
+            if isinstance(response, dict):
+                response_display = json.dumps(response)
+            else:
+                response_display = response.model_dump_json()
+            return [
+                Message(role=Role.USER, content=input_text),
+                Message(role=Role.ASSISTANT, content=response_display),
+            ]
+        else:
+            return Message(role=Role.USER, content=input_text)
+
+    messages = [
+        Message(
+            role=Role.SYSTEM,
+            content="Parse the following text with the provided JSON schema.",
+        ),
+    ]
+    if examples is not None:
+        for example in examples:
+            example_msgs = display_sample(t=example.text, ty=typ, response=example.response)
+            assert isinstance(example_msgs, list) and all(
+                isinstance(msg, Message) for msg in example_msgs
+            ), example_msgs
+            messages.extend(example_msgs)
+    messages.append(display_sample(t=text, ty=typ))
+    return messages
+
+
+def build_classify_state(
+    instruction: str,
+    text: str,
+    options: list[str],
+    examples: list[ClassifyDatapoint] | None = None,
+    suffix_strategy: PromptSuffixStrategy = PromptSuffixStrategy.JSON,
+) -> tuple[list[Message], dict[str, int]]:
+    def display_sample(
+        instr: str, t: str, opts: list[str], response: int | None = None
+    ) -> list[Message] | tuple[Message, dict[str, int]]:
+        choices_display, decode_map = display_choices(opts)
+        input_text = force_json_prompt(
+            f"Instruction:\n{instr}\n\nText:\n{t}\n\nChoices:\n{choices_display}",
+            suffix_strategy=suffix_strategy,
+        )
+        if response is not None:
+            response_label = None
+            for label, idx in decode_map.items():
+                if idx == response:
+                    response_label = label
+                    break
+            assert response_label is not None, f"Invalid response: {response}"
+            return [
+                Message(role=Role.USER, content=input_text),
+                Message(
+                    role=Role.ASSISTANT,
+                    content=f'{{"classification": "{response_label}"}}',
+                ),
+            ]
+        else:
+            return Message(role=Role.USER, content=input_text), decode_map
+
+    messages = [
+        Message(
+            role=Role.SYSTEM,
+            content='Classify the following text with the provided instruction and choices. To classify, provide the key of the choice:\n{"classification": string}\n\nFor example, if the correct choice is \'Z. description of choice Z\', then provide \'Z\' as the classification as valid JSON:\n{"classification": "Z"}',
+        ),
+    ]
+    if examples is not None:
+        for example in examples:
+            example_msgs = display_sample(
+                instr=example.instruction,
+                t=example.text,
+                opts=example.options,
+                response=example.response,
+            )
+            assert isinstance(example_msgs, list) and all(
+                isinstance(msg, Message) for msg in example_msgs
+            ), example_msgs
+            messages.extend(example_msgs)
+    message, decode_map = display_sample(instr=instruction, t=text, opts=options)
+    messages.append(message)
+    return messages, decode_map
+
+
+class ChatModel(GeneralModel):
+    @abc.abstractmethod
+    def generate_message(
+        self, messages: list[Message], force_json: bool, temperature: float | None = None
+    ) -> Message:
+        raise NotImplementedError
+
+    def handle_generate_message_response(
+        self, prompt: list[dict[str, str] | Message], content: str, force_json: bool
+    ) -> Message:
+        if force_json:
+            try:
+                parsed = parse_json_or_json_markdown(content)
+            except (json.JSONDecodeError, ValueError) as e:
+                msgs = []
+                for msg in prompt:
+                    if isinstance(msg, Message):
+                        msgs.append(msg.model_dump())
+                    else:
+                        msgs.append(msg)
+                raise ModelError(
+                    short_message=f"Failed to parse JSON: {content}",
+                    prompt=msgs,
+                    response=content,
+                ) from e
+            cleaned = clean_top_level_keys(parsed)
+            return Message(role=Role.ASSISTANT, content=content, obj=cleaned)
+        return Message(role=Role.ASSISTANT, content=content, obj=None)
+
+    def build_generate_message_state(self, messages: list[Message]) -> list[dict[str, str]]:
+        msgs: list[dict[str, str]] = []
+        for msg in messages:
+            if msg.obj is not None:
+                content = json.dumps(msg.obj)
+            else:
+                content = msg.content
+            msgs.append({"role": msg.role.value, "content": content})
+        return msgs
+
+    def _handle_classify_response(self, res: Message, decode_map: dict[str, int]) -> int:
+        assert res.obj is not None
+        if "classification" not in res.obj:
+            raise ModelError(f"Invalid response from model: {res.content}")
+        choice = res.obj["classification"]
+        if choice not in decode_map:
+            key = try_classify_recover(s=choice, decode_map=decode_map)
+            if key is not None:
+                return decode_map[key]
+            raise ModelError(f"Invalid choice: {choice}")
+        return decode_map[choice]
+
+    def classify(
+        self,
+        instruction: str,
+        text: str,
+        options: list[str],
+        examples: list[ClassifyDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> int:
+        messages, decode_map = build_classify_state(instruction, text, options, examples=examples)
+        res = self.generate_message(messages, force_json=True, temperature=temperature)
+        return self._handle_classify_response(res, decode_map)
+
+    def parse(
+        self,
+        text: str,
+        typ: type[T] | dict[str, Any],
+        examples: list[ParseDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> T | PartialObj | dict[str, Any]:
+        messages = build_parse_state(text, typ, examples=examples)
+        res = self.generate_message(messages, force_json=True, temperature=temperature)
+        assert res.obj is not None
+        return json_response_to_obj_or_partial_obj(response=res.obj, typ=typ)
+
+    def generate(
+        self,
+        instruction: str,
+        text: str,
+        examples: list[GenerateDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> str:
+        messages = build_generate_state(instruction=instruction, text=text, examples=examples)
+        return self.generate_message(messages, force_json=False, temperature=temperature).content
+
+    def _handle_parse_force_response(
+        self, res: Message, typ: type[T] | dict[str, Any]
+    ) -> T | dict[str, Any]:
+        assert res.obj is not None
+        obj = json_response_to_obj_or_partial_obj(response=res.obj, typ=typ)
+        if not isinstance(typ, dict) and isinstance(obj, dict):
+            raise ModelError(f"Invalid response from model: {res.content}")
+        return obj
+
+    def parse_force(
+        self,
+        instruction: str,
+        typ: type[T] | dict[str, Any],
+        text: str | None = None,
+        examples: list[ParseForceDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> T | dict[str, Any]:
+        messages = build_parse_force_state(
+            instruction=instruction,
+            typ=typ,
+            text=text,
+            examples=examples,
+        )
+        res = self.generate_message(messages, force_json=True, temperature=temperature)
+        return self._handle_parse_force_response(res, typ)
+
+    def _handle_score_response(
+        self,
+        res: Message,
+        min: int,
+        max: int,
+    ) -> int:
+        if res.obj is None or "score" not in res.obj:
+            raise ModelError(f"Invalid response from model: {res.content}")
+        score = res.obj["score"]
+        if not isinstance(score, int):
+            raise ModelError(f"Invalid score type: {type(score)}")
+        if score < min or score > max:
+            raise ModelError(f"Invalid score value: {score}")
+        return score
+
+    def score(
+        self,
+        instruction: str,
+        text: str,
+        min: int,
+        max: int,
+        examples: list[ScoreDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> int:
+        messages = build_score_state(instruction, text, min, max, examples=examples)
+        res = self.generate_message(messages, force_json=True, temperature=temperature)
+        return self._handle_score_response(res, min, max)
+
+
+def build_prompts(
+    dps: list[Datapoint], prompt_suffix_strategy: PromptSuffixStrategy | None
+) -> list[str | list[Message]]:
+    if len(dps) == 0:
+        return []
+    typ = type(dps[0])
+    for i, dp in enumerate(dps):
+        if not isinstance(dp, typ):
+            raise ValueError(
+                f"All elements must be of type Datapoint, expected type {typ} at index {i}, got {type(dp)}"
+            )
+    if isinstance(dps[0], ParseDatapoint):
+        build_func = build_parse_prompts
+    elif isinstance(dps[0], BinaryClassifyDatapoint):
+        build_func = build_binary_classify_prompts
+    elif isinstance(dps[0], ClassifyDatapoint):
+        build_func = build_classify_prompts
+    elif isinstance(dps[0], ParseForceDatapoint):
+        build_func = build_parse_force_prompts
+    elif isinstance(dps[0], GenerateDatapoint):
+        build_func = build_generate_prompts
+    elif isinstance(dps[0], ScoreDatapoint):
+        build_func = build_score_prompts
+    else:
+        raise ValueError(f"Unknown datapoint type: {type(dps[0])}")
+    return build_func(dps, suffix_strategy=prompt_suffix_strategy)
+
+
+def build_parse_prompts(
+    dps: list[ParseDatapoint],
+    suffix_strategy: PromptSuffixStrategy | None = None,
+) -> list[str | list[Message]]:
+    datapoints = []
+    for dp in dps:
+        json_response_object = (
+            dp.response.model_dump_json()
+            if isinstance(dp.response, BaseModel)
+            else json.dumps(dp.response)
+        )
+        prompt_msgs = build_parse_state(
+            text=dp.text,
+            typ=dp.typ,
+            suffix_strategy=(
+                suffix_strategy if suffix_strategy is not None else PromptSuffixStrategy.JSON
+            ),
+        )
+        json_response = apply_suffix_strategy(
+            response=json_response_object, suffix_strategy=suffix_strategy
+        )
+        datapoints.append(prompt_msgs + [Message(role=Role.ASSISTANT, content=json_response)])
+    return datapoints
+
+
+def build_binary_classify_prompts(
+    dps: list[BinaryClassifyDatapoint],
+    suffix_strategy: PromptSuffixStrategy | None = None,
+) -> list[str | list[Message]]:
+    return build_classify_prompts(
+        [
+            ClassifyDatapoint(
+                instruction=dp.instruction,
+                text=dp.text,
+                options=["true", "false"],
+                response=0 if dp.response else 1,
+            )
+            for dp in dps
+        ],
+        suffix_strategy=suffix_strategy,
+    )
+
+
+def build_classify_prompts(
+    dps: list[ClassifyDatapoint],
+    suffix_strategy: PromptSuffixStrategy | None = None,
+) -> list[str | list[Message]]:
+    def label_idx_to_label_json(idx: int, decode_map: dict[str, int]) -> str:
+        label = None
+        for k, v in decode_map.items():
+            if v == idx:
+                label = k
+                break
+        if label is None:
+            raise ValueError(f"Label index {idx} not found in decode map")
+        return f'{{"classification": "{label}"}}'
+
+    datapoints = []
+    for dp in dps:
+        suffix_strategy = PromptSuffixStrategy.JSON if suffix_strategy is None else suffix_strategy
+        prompt_msgs, decode_map = build_classify_state(
+            instruction=dp.instruction,
+            text=dp.text,
+            options=dp.options,
+            suffix_strategy=suffix_strategy,
+        )
+        json_response_object = label_idx_to_label_json(idx=dp.response, decode_map=decode_map)
+        json_response = apply_suffix_strategy(
+            response=json_response_object, suffix_strategy=suffix_strategy
+        )
+        datapoints.append(
+            prompt_msgs
+            + [
+                Message(
+                    role=Role.ASSISTANT,
+                    content=json_response,
+                )
+            ]
+        )
+    return datapoints
+
+
+def build_parse_force_prompts(
+    dps: list[ParseForceDatapoint],
+    suffix_strategy: PromptSuffixStrategy | None = None,
+) -> list[str | list[Message]]:
+    datapoints = []
+    for dp in dps:
+        json_response_obj = (
+            dp.response.model_dump_json()
+            if isinstance(dp.response, BaseModel)
+            else json.dumps(dp.response)
+        )
+        suffix_strategy = PromptSuffixStrategy.JSON if suffix_strategy is None else suffix_strategy
+        prompt_msgs = build_parse_force_state(
+            instruction=dp.instruction,
+            text=dp.text,
+            typ=dp.typ,
+            suffix_strategy=suffix_strategy,
+        )
+        json_response = apply_suffix_strategy(
+            response=json_response_obj, suffix_strategy=suffix_strategy
+        )
+        datapoints.append(prompt_msgs + [Message(role=Role.ASSISTANT, content=json_response)])
+    return datapoints
+
+
+def build_generate_prompts(dps: list[GenerateDatapoint]) -> list[str | list[Message]]:
+    datapoints = []
+    for dp in dps:
+        prompt_msgs = build_generate_state(instruction=dp.instruction, text=dp.text)
+        datapoints.append(prompt_msgs + [Message(role=Role.ASSISTANT, content=dp.response)])
+    return datapoints
+
+
+def build_score_prompts(
+    dps: list[ScoreDatapoint],
+    suffix_strategy: PromptSuffixStrategy | None = None,
+) -> list[str | list[Message]]:
+    datapoints = []
+    for dp in dps:
+        json_response_object = f'{{"score": {dp.response}}}'
+        suffix_strategy = (
+            suffix_strategy if suffix_strategy is not None else PromptSuffixStrategy.JSON
+        )
+        prompt_msgs = build_score_state(
+            instruction=dp.instruction,
+            text=dp.text,
+            min=dp.min,
+            max=dp.max,
+            suffix_strategy=suffix_strategy,
+        )
+        json_response = apply_suffix_strategy(
+            response=json_response_object, suffix_strategy=suffix_strategy
+        )
+        datapoints.append(prompt_msgs + [Message(role=Role.ASSISTANT, content=json_response)])
+    return datapoints
+
+
+def apply_suffix_strategy(response: str, suffix_strategy: PromptSuffixStrategy) -> str:
+    if suffix_strategy == PromptSuffixStrategy.JSON:
+        return response
+    elif suffix_strategy == PromptSuffixStrategy.JSON_MD_BLOCK:
+        return add_md_tag(response)
+    else:
+        raise ValueError(f"Unknown suffix strategy: {suffix_strategy}")

chapter2/prompt-engineering/tau_bench/model_utils/model/claude.py

@@ -0,0 +1,138 @@
+import json
+import os
+
+from tau_bench.model_utils.api.datapoint import Datapoint
+from tau_bench.model_utils.model.chat import ChatModel, Message
+from tau_bench.model_utils.model.completion import approx_cost_for_datapoint, approx_prompt_str
+from tau_bench.model_utils.model.general_model import wrap_temperature
+from tau_bench.model_utils.model.utils import approx_num_tokens
+
+DEFAULT_CLAUDE_MODEL = "claude-3-5-sonnet-20240620"
+DEFAULT_MAX_TOKENS = 8192
+ENV_VAR_API_KEY = "ANTHROPIC_API_KEY"
+
+PRICE_PER_INPUT_TOKEN_MAP = {
+    "claude-3-5-sonnet-20240620": 3 / 1000000,
+}
+INPUT_PRICE_PER_TOKEN_FALLBACK = 15 / 1000000
+
+CAPABILITY_SCORE_MAP = {
+    "claude-3-5-sonnet-20240620": 1.0,
+}
+CAPABILITY_SCORE_FALLBACK = 0.5
+
+# TODO: implement
+LATENCY_MS_PER_OUTPUT_TOKEN_MAP = {}
+# TODO: implement
+LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK = 0.0
+
+MAX_CONTEXT_LENGTH_MAP = {
+    "claude-3-5-sonnet-20240620": 8192,
+}
+MAX_CONTEXT_LENGTH_FALLBACK = 8192
+
+
+class ClaudeModel(ChatModel):
+    def __init__(
+        self,
+        model: str | None = None,
+        api_key: str | None = None,
+        temperature: float = 0.0,
+    ) -> None:
+        from anthropic import Anthropic, AsyncAnthropic
+
+        if model is None:
+            self.model = DEFAULT_CLAUDE_MODEL
+        else:
+            self.model = model
+
+        api_key = None
+        if api_key is None:
+            api_key = os.getenv(ENV_VAR_API_KEY)
+            if api_key is None:
+                raise ValueError(f"{ENV_VAR_API_KEY} environment variable is not set")
+        # `anthropic-beta` header is needed for the 8192 context length (https://docs.anthropic.com/en/docs/about-claude/models)
+        self.client = Anthropic(
+            api_key=api_key, default_headers={"anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15"}
+        )
+        self.async_client = AsyncAnthropic(api_key=api_key)
+        self.temperature = temperature
+
+    def get_approx_cost(self, dp: Datapoint) -> float:
+        cost_per_token = PRICE_PER_INPUT_TOKEN_MAP.get(self.model, INPUT_PRICE_PER_TOKEN_FALLBACK)
+        return approx_cost_for_datapoint(dp=dp, price_per_input_token=cost_per_token)
+
+    def get_latency(self, dp: Datapoint) -> float:
+        latency_per_output_token = LATENCY_MS_PER_OUTPUT_TOKEN_MAP.get(
+            self.model, LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK
+        )
+        return approx_cost_for_datapoint(dp=dp, price_per_input_token=latency_per_output_token)
+
+    def get_capability(self) -> float:
+        return CAPABILITY_SCORE_MAP.get(self.model, CAPABILITY_SCORE_FALLBACK)
+
+    def supports_dp(self, dp: Datapoint) -> bool:
+        prompt = approx_prompt_str(dp)
+        return approx_num_tokens(prompt) <= MAX_CONTEXT_LENGTH_MAP.get(
+            self.model, MAX_CONTEXT_LENGTH_FALLBACK
+        )
+
+    def _remap_messages(self, messages: list[dict[str, str]]) -> list[dict[str, str]]:
+        remapped: list[dict[str, str]] = []
+        is_user = True
+        for i, message in enumerate(messages):
+            role = message["role"]
+            if role == "assistant":
+                if i == 0:
+                    raise ValueError(
+                        f"First message must be a system or user message, got {[m['role'] for m in messages]}"
+                    )
+                elif is_user:
+                    raise ValueError(
+                        f"Must alternate between user and assistant, got {[m['role'] for m in messages]}"
+                    )
+                remapped.append(message)
+                is_user = True
+            else:
+                if is_user:
+                    remapped.append({"role": "user", "content": message["content"]})
+                    is_user = False
+                else:
+                    if remapped[-1]["role"] != "user":
+                        raise ValueError(
+                            f"Invalid sequence, expected user message but got {[m['role'] for m in messages]}"
+                        )
+                    remapped[-1]["content"] += "\n\n" + message["content"]
+        return remapped
+
+    def build_generate_message_state(
+        self,
+        messages: list[Message],
+    ) -> list[dict[str, str]]:
+        msgs: list[dict[str, str]] = []
+        for msg in messages:
+            if msg.obj is not None:
+                content = json.dumps(msg.obj)
+            else:
+                content = msg.content
+            msgs.append({"role": msg.role.value, "content": content})
+        return self._remap_messages(msgs)
+
+    def generate_message(
+        self,
+        messages: list[Message],
+        force_json: bool,
+        temperature: float | None = None,
+    ) -> Message:
+        if temperature is None:
+            temperature = self.temperature
+        msgs = self.build_generate_message_state(messages)
+        res = self.client.messages.create(
+            model=self.model,
+            messages=msgs,
+            temperature=wrap_temperature(temperature),
+            max_tokens=DEFAULT_MAX_TOKENS,
+        )
+        return self.handle_generate_message_response(
+            prompt=msgs, content=res.content[0].text, force_json=force_json
+        )

chapter2/prompt-engineering/tau_bench/model_utils/model/completion.py

@@ -0,0 +1,538 @@
+import abc
+import json
+from typing import Any, TypeVar
+
+from pydantic import BaseModel
+
+from tau_bench.model_utils.api.datapoint import (
+    BinaryClassifyDatapoint,
+    ClassifyDatapoint,
+    Datapoint,
+    GenerateDatapoint,
+    ParseDatapoint,
+    ParseForceDatapoint,
+    ScoreDatapoint,
+)
+from tau_bench.model_utils.api.types import PartialObj
+from tau_bench.model_utils.model.exception import ModelError
+from tau_bench.model_utils.model.general_model import GeneralModel
+from tau_bench.model_utils.model.utils import (
+    add_md_close_tag,
+    approx_num_tokens,
+    display_choices,
+    json_response_to_obj_or_partial_obj,
+    optionalize_type,
+    parse_json_or_json_markdown,
+    try_classify_recover,
+    type_to_json_schema_string,
+)
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class Score(BaseModel):
+    score: int
+
+
+class Classification(BaseModel):
+    classification: str
+
+
+def task_prompt(task: str, text: str) -> str:
+    return f"# Task\n{task}\n\n{text}"
+
+
+def force_json_prompt(text: str, with_prefix: bool = False) -> str:
+    suffix = (
+        'For example:\nassistant:```json\n{"key": "value"}\n```'
+        if not with_prefix
+        else "\n\n```json\n"
+    )
+    return f"{text}\n\nThe result should be a valid JSON object in a markdown block only. {suffix}"
+
+
+def build_score_state(
+    instruction: str,
+    text: str,
+    min: int,
+    max: int,
+    examples: list[ScoreDatapoint] | None = None,
+) -> str:
+    def display_sample(instr: str, t: str, min: int, max: int, response: int | None = None) -> str:
+        p = task_prompt(
+            task='Score the following text with the provided instruction and range as an integer value in valid JSON:\n{"score": number}',
+            text=force_json_prompt(
+                f"Instruction:\n{instr}\n\nText:\n{t}\n\nRange:\n[{min}, {max}]",
+                with_prefix=True,
+            ),
+        )
+        if response is not None:
+            # the json markdown block is opened in the prompt
+            return f'{p}\n{{"score": {response}}}\n```'
+        return p
+
+    p = (
+        "\n\n".join(
+            [display_sample(ex.instruction, ex.text, min, max, ex.response) for ex in examples]
+        )
+        if examples is not None
+        else ""
+    )
+    return f"{p}\n\n{display_sample(instr=instruction, t=text, min=min, max=max)}"
+
+
+def build_parse_force_state(
+    instruction: str,
+    typ: type[T] | dict[str, Any],
+    text: str | None = None,
+    examples: list[ParseForceDatapoint] | None = None,
+) -> str:
+    def display_sample(
+        instr: str,
+        t: str,
+        ty: type[T] | dict[str, Any],
+        response: T | dict[str, Any] | None = None,
+    ) -> str:
+        if isinstance(ty, dict):
+            json_schema_string = json.dumps(ty)
+        else:
+            json_schema_string = type_to_json_schema_string(ty)
+        text_insert = "" if t is None else f"\n\nText:\n{t}"
+        input_text = force_json_prompt(
+            text=f"Instruction:\n{instr}{text_insert}\n\nSchema:\n{json_schema_string}",
+            with_prefix=True,
+        )
+        if response is not None:
+            if isinstance(response, dict):
+                response_display = json.dumps(response)
+            else:
+                response_display = response.model_dump_json()
+            # the json markdown block is opened in the prompt
+            return f"{input_text}\n{response_display}\n```"
+        return input_text
+
+    p = (
+        "".join(
+            [
+                display_sample(
+                    instr=ex.instruction,
+                    t=ex.text,
+                    ty=ex.typ,
+                    response=ex.response,
+                )
+                for ex in examples
+            ]
+        )
+        + "\n\n"
+        if examples is not None and len(examples) > 0
+        else ""
+    )
+    p += display_sample(instr=instruction, t=text, ty=typ)
+    return task_prompt(
+        task="Generate an object with the provided instruction, text, and schema.",
+        text=p,
+    )
+
+
+def build_parse_state(
+    text: str,
+    typ: type[T] | dict[str, Any],
+    examples: list[ParseDatapoint] | None = None,
+) -> str:
+    instruction = "Parse the following text with the provided JSON schema."
+
+    def display_sample(
+        t: str,
+        ty: type[T] | dict[str, Any],
+        response: T | PartialObj | dict[str, Any] | None = None,
+    ) -> str:
+        if isinstance(ty, dict):
+            json_schema_string = json.dumps(ty)
+        else:
+            optionalized_typ = optionalize_type(ty)
+            json_schema_string = type_to_json_schema_string(optionalized_typ)
+        # instruction is repeated to emphasize the task
+        prompt = task_prompt(
+            task=instruction,
+            text=force_json_prompt(
+                f"Text:\n{t}\n\nSchema:\n{json_schema_string}", with_prefix=True
+            ),
+        )
+        if response is None:
+            return prompt
+        if isinstance(response, dict):
+            response_display = json.dumps(response)
+        else:
+            response_display = response.model_dump_json()
+        # the json markdown block is opened in the prompt
+        json_response = f"{response_display}\n```"
+        return f"{prompt}\n{json_response}"
+
+    p = ""
+    if examples is not None and len(examples) > 0:
+        p = "\n\n".join(
+            [display_sample(t=ex.text, ty=ex.typ, response=ex.response) for ex in examples]
+        )
+    return f"{p}\n\n{display_sample(t=text, ty=typ)}"
+
+
+def build_classify_state(
+    instruction: str,
+    text: str,
+    options: list[str],
+    examples: list[ClassifyDatapoint] | None = None,
+) -> tuple[str, dict[str, int]]:
+    def display_sample(
+        instr: str, t: str, opts: list[str], response: int | None = None
+    ) -> str | tuple[str, dict[str, int]]:
+        choices_display, decode_map = display_choices(opts)
+        input_text = force_json_prompt(
+            f"Instruction:\n{instr}\n\nText:\n{t}\n\nChoices:\n{choices_display}",
+            with_prefix=True,
+        )
+        prompt = task_prompt(task=instr, text=input_text)
+        if response is not None:
+            label = None
+            for k, v in decode_map.items():
+                if v == response:
+                    label = k
+                    break
+            assert label is not None
+            # the json markdown block is opened in the prompt
+            json_display = f'{{"classification": "{label}"}}\n```'
+            return f"{prompt}\n{json_display}"
+        return prompt, decode_map
+
+    p = 'Classify the following text with the provided instruction and choices. To classify, provide the key of the choice:\n{"classification": string}\n\nFor example, if the correct choice is \'Z. description of choice Z\', then provide \'Z\' as the classification as valid JSON:\n```json\n{"classification": "Z"}\n```'
+    if examples is not None and len(examples) > 0:
+        example_displays = "\n\n".join(
+            [
+                display_sample(
+                    instr=ex.instruction,
+                    t=ex.text,
+                    opts=ex.options,
+                    response=ex.response,
+                )
+                for ex in examples
+            ]
+        )
+        p += f"\n\n{example_displays}"
+    prompt, decode_map = display_sample(instr=instruction, t=text, opts=options)
+    return f"{p}\n\n{prompt}", decode_map
+
+
+def build_generate_state(
+    instruction: str,
+    text: str,
+    examples: list[GenerateDatapoint] | None = None,
+) -> str:
+    def display_sample(instr: str, t: str, response: str | None = None) -> str:
+        prompt = task_prompt(task=instr, text=t)
+        if response is not None:
+            return f"{prompt}\n\nText: {response}"
+        return prompt
+
+    prompt = (
+        "\n\n".join([display_sample(ex.instruction, ex.text) for ex in examples]) + "\n\n"
+        if examples is not None and len(examples) > 0
+        else ""
+    )
+    return f"{prompt}\n\n{display_sample(instruction, text)}\n\nText:"
+
+
+class CompletionModel(GeneralModel):
+    @abc.abstractmethod
+    def generate_from_prompt(self, prompt: str, temperature: float | None = None) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def parse_force_from_prompt(
+        self, prompt: str, typ: BaseModel | dict[str, Any], temperature: float | None = None
+    ) -> dict[str, Any]:
+        raise NotImplementedError
+
+    def handle_parse_force_response(self, prompt: str, content: str) -> dict[str, Any]:
+        try:
+            return parse_json_or_json_markdown(content)
+        except (json.decoder.JSONDecodeError, ValueError) as e:
+            raise ModelError(
+                short_message=f"Failed to decode JSON: {content}", prompt=prompt, response=content
+            ) from e
+
+    def _handle_classify_response(self, res: dict[str, int], decode_map: dict[str, int]) -> int:
+        if "classification" not in res:
+            raise ModelError(f"Invalid response from model: {res}")
+        choice = res["classification"]
+        if choice not in decode_map.keys():
+            key = try_classify_recover(s=choice, decode_map=decode_map)
+            if key is not None:
+                return decode_map[key]
+            raise ModelError(f"Invalid choice: {choice}")
+        return decode_map[choice]
+
+    def classify(
+        self,
+        instruction: str,
+        text: str,
+        options: list[str],
+        examples: list[ClassifyDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> int:
+        prompt, decode_map = build_classify_state(instruction, text, options, examples=examples)
+        res = self.parse_force_from_prompt(prompt, typ=Classification, temperature=temperature)
+        return self._handle_classify_response(res, decode_map)
+
+    def parse(
+        self,
+        text: str,
+        typ: type[T] | dict[str, Any],
+        examples: list[ParseDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> T | PartialObj | dict[str, Any]:
+        prompt = build_parse_state(text, typ, examples=examples)
+        res = self.parse_force_from_prompt(prompt=prompt, typ=typ, temperature=temperature)
+        return json_response_to_obj_or_partial_obj(response=res, typ=typ)
+
+    def generate(
+        self,
+        instruction: str,
+        text: str,
+        examples: list[GenerateDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> str:
+        prompt = build_generate_state(instruction=instruction, text=text, examples=examples)
+        return self.generate_from_prompt(prompt=prompt, temperature=temperature)
+
+    def _handle_parse_force_response(self, res: dict[str, Any], typ: type[T]) -> T:
+        obj = json_response_to_obj_or_partial_obj(response=res, typ=typ)
+        if isinstance(obj, dict):
+            raise ModelError(f"Invalid response from model: {res}")
+        return obj
+
+    def parse_force(
+        self,
+        instruction: str,
+        typ: type[T] | dict[str, Any],
+        text: str | None = None,
+        examples: list[ParseForceDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> T | dict[str, Any]:
+        prompt = build_parse_force_state(
+            instruction=instruction, text=text, typ=typ, examples=examples
+        )
+        res = self.parse_force_from_prompt(prompt=prompt, typ=typ, temperature=temperature)
+        return self._handle_parse_force_response(res, typ)
+
+    def _handle_score_response(
+        self,
+        res: dict[str, Any],
+        min: int,
+        max: int,
+    ) -> int:
+        if res is None or "score" not in res:
+            raise ModelError(f"Invalid response from model: {res}")
+        score = res["score"]
+        if not isinstance(score, int):
+            raise ModelError(f"Invalid score type: {type(score)}")
+        if score < min or score > max:
+            raise ModelError(f"Invalid score value: {score}")
+        return score
+
+    def score(
+        self,
+        instruction: str,
+        text: str,
+        min: int,
+        max: int,
+        examples: list[ScoreDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> int:
+        prompt = build_score_state(instruction, text, min, max, examples=examples)
+        res = self.parse_force_from_prompt(prompt=prompt, typ=Score, temperature=temperature)
+        return self._handle_score_response(res, min, max)
+
+
+def build_prompts(dps: list[Datapoint], include_response: bool = True) -> list[str]:
+    if len(dps) == 0:
+        return []
+    typ = type(dps[0])
+    for i, dp in enumerate(dps):
+        if not isinstance(dp, typ):
+            raise ValueError(
+                f"All elements must be of type Datapoint, expected type {typ} at index {i}, got {type(dp)}"
+            )
+    if isinstance(dps[0], ParseDatapoint):
+        build_func = build_parse_prompts
+    elif isinstance(dps[0], BinaryClassifyDatapoint):
+        build_func = build_binary_classify_prompts
+    elif isinstance(dps[0], ClassifyDatapoint):
+        build_func = build_classify_prompts
+    elif isinstance(dps[0], ParseForceDatapoint):
+        build_func = build_parse_force_prompts
+    elif isinstance(dps[0], GenerateDatapoint):
+        build_func = build_generate_prompts
+    elif isinstance(dps[0], ScoreDatapoint):
+        build_func = build_score_prompts
+    else:
+        raise ValueError(f"Unknown datapoint type: {type(dps[0])}")
+    return build_func(dps, include_response)
+
+
+def build_parse_prompts(
+    dps: list[ParseDatapoint],
+    include_response: bool = True,
+) -> list[str]:
+    datapoints = []
+    for dp in dps:
+        json_response_object = (
+            dp.response.model_dump_json()
+            if isinstance(dp.response, BaseModel)
+            else json.dumps(dp.response)
+        )
+        prompt = build_parse_state(text=dp.text, typ=dp.typ)
+        if include_response:
+            json_response = add_md_close_tag(json_response_object)
+            datapoints.append(prompt + json_response)
+        else:
+            datapoints.append(prompt)
+    return datapoints
+
+
+def build_binary_classify_prompts(
+    dps: list[BinaryClassifyDatapoint],
+    include_response: bool = True,
+) -> list[str]:
+    return build_classify_prompts(
+        [
+            ClassifyDatapoint(
+                instruction=dp.instruction,
+                text=dp.text,
+                options=["true", "false"],
+                response=0 if dp.response else 1,
+            )
+            for dp in dps
+        ],
+        include_response=include_response,
+    )
+
+
+def build_classify_prompts(
+    dps: list[ClassifyDatapoint],
+    include_response: bool = True,
+) -> list[str]:
+    def label_idx_to_label_json(idx: int, decode_map: dict[str, int]) -> str:
+        label = None
+        for k, v in decode_map.items():
+            if v == idx:
+                label = k
+                break
+        if label is None:
+            raise ValueError(f"Label index {idx} not found in decode map")
+        return f'{{"classification": "{label}"}}'
+
+    datapoints = []
+    for dp in dps:
+        prompt, decode_map = build_classify_state(
+            instruction=dp.instruction, text=dp.text, options=dp.options
+        )
+        if include_response:
+            json_response_object = label_idx_to_label_json(idx=dp.response, decode_map=decode_map)
+            json_response = add_md_close_tag(json_response_object)
+            datapoints.append(prompt + json_response)
+        else:
+            datapoints.append(prompt)
+    return datapoints
+
+
+def build_parse_force_prompts(
+    dps: list[ParseForceDatapoint],
+    include_response: bool = True,
+) -> list[str]:
+    datapoints = []
+    for dp in dps:
+        json_response_obj = (
+            dp.response.model_dump_json()
+            if isinstance(dp.response, BaseModel)
+            else json.dumps(dp.response)
+        )
+        prompt = build_parse_force_state(
+            instruction=dp.instruction,
+            text=dp.text,
+            typ=dp.typ,
+        )
+        if include_response:
+            json_response = add_md_close_tag(json_response_obj)
+            datapoints.append(prompt + json_response)
+        else:
+            datapoints.append(prompt)
+    return datapoints
+
+
+def build_generate_prompts(
+    dps: list[GenerateDatapoint], include_response: bool = True
+) -> list[str]:
+    datapoints = []
+    for dp in dps:
+        prompt = build_generate_state(instruction=dp.instruction, text=dp.text)
+        if include_response:
+            datapoints.append(prompt + dp.response)
+        else:
+            datapoints.append(prompt)
+    return datapoints
+
+
+def build_score_prompts(
+    dps: list[ScoreDatapoint],
+    include_response: bool = True,
+) -> list[str]:
+    datapoints = []
+    for dp in dps:
+        json_response_object = f'{{"score": {dp.response}}}'
+        prompt = build_score_state(
+            instruction=dp.instruction,
+            text=dp.text,
+            min=dp.min,
+            max=dp.max,
+        )
+        if include_response:
+            json_response = add_md_close_tag(json_response_object)
+            datapoints.append(prompt + json_response)
+        else:
+            datapoints.append(prompt)
+    return datapoints
+
+
+# TODO: handle examples
+def approx_prompt_str(dp: Datapoint, include_response: bool = False) -> str:
+    return build_prompts(dps=[dp], include_response=include_response)[0]
+
+
+# TODO: handle examples
+def approx_cost_for_datapoint(
+    dp: Datapoint,
+    price_per_input_token: float,
+) -> float:
+    """For now, we approximate the cost of a datapoint as the cost of the input (output tokens are priced as input tokens as well)."""
+    prompt = approx_prompt_str(dp, include_response=True)
+    assert isinstance(prompt, str)
+    return price_per_input_token * approx_num_tokens(prompt)
+
+
+# TODO: handle examples
+def approx_latency_for_datapoint(dp: Datapoint, latency_ms_per_output_token: float) -> float:
+    if isinstance(dp, BinaryClassifyDatapoint) or isinstance(dp, ClassifyDatapoint):
+        approx_response = '{"classification": 0}'
+    elif isinstance(dp, ParseDatapoint):
+        # this is extremely approximate
+        approx_response = '{"street": "main st", "city": "san francisco", "state": "CA"}'
+    elif isinstance(dp, GenerateDatapoint):
+        # this is extremely approximate
+        approx_response = "This is a generated text response."
+    elif isinstance(dp, ParseForceDatapoint):
+        # this is extremely approximate
+        approx_response = '{"street": "main st", "city": "san francisco", "state": "CA"}'
+    elif isinstance(dp, ScoreDatapoint):
+        approx_response = '{"score": 0}'
+    else:
+        raise ValueError(f"Unsupported datapoint type: {type(dp)}")
+    return latency_ms_per_output_token * approx_num_tokens(approx_response)

chapter2/prompt-engineering/tau_bench/model_utils/model/exception.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+T = TypeVar("T")
+
+
+class ModelError(Exception):
+    def __init__(
+        self,
+        short_message: str,
+        prompt: str | list[dict[str, str]] | None = None,
+        response: str | None = None,
+    ) -> None:
+        super().__init__(short_message)
+        self.short_message = short_message
+        self.prompt = prompt
+        self.response = response
+
+
+@dataclass
+class Result(Generic[T]):
+    value: T | None
+    error: ModelError | None

chapter2/prompt-engineering/tau_bench/model_utils/model/general_model.py

@@ -0,0 +1,187 @@
+import abc
+from typing import Any, TypeVar
+
+from pydantic import BaseModel
+
+from tau_bench.model_utils.api.datapoint import (
+    BinaryClassifyDatapoint,
+    ClassifyDatapoint,
+    GenerateDatapoint,
+    ParseDatapoint,
+    ParseForceDatapoint,
+    ScoreDatapoint,
+)
+from tau_bench.model_utils.api.types import PartialObj
+from tau_bench.model_utils.model.model import (
+    BinaryClassifyModel,
+    ClassifyModel,
+    GenerateModel,
+    ParseForceModel,
+    ParseModel,
+    Platform,
+    ScoreModel,
+)
+
+T = TypeVar("T", bound=BaseModel)
+
+LLM_SAMPLING_TEMPERATURE_EPS = 1e-5
+
+
+def wrap_temperature(temperature: float) -> float:
+    return max(temperature, LLM_SAMPLING_TEMPERATURE_EPS)
+
+
+class GeneralModel(
+    ClassifyModel,
+    BinaryClassifyModel,
+    ParseModel,
+    GenerateModel,
+    ParseForceModel,
+    ScoreModel,
+):
+    @abc.abstractmethod
+    def classify(
+        self,
+        instruction: str,
+        text: str,
+        options: list[str],
+        examples: list[ClassifyDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> int:
+        raise NotImplementedError
+
+    def binary_classify(
+        self,
+        instruction: str,
+        text: str,
+        examples: list[BinaryClassifyDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> bool:
+        return (
+            self.classify(
+                instruction,
+                text,
+                ["true", "false"],
+                examples=(
+                    None
+                    if examples is None
+                    else [
+                        ClassifyDatapoint(
+                            instruction=example.instruction,
+                            text=example.text,
+                            options=["true", "false"],
+                            response=0 if example.response else 1,
+                        )
+                        for example in examples
+                    ]
+                ),
+                temperature=temperature,
+            )
+            == 0
+        )
+
+    @abc.abstractmethod
+    def parse(
+        self,
+        text: str,
+        typ: type[T] | dict[str, Any],
+        examples: list[ParseDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> T | PartialObj | dict[str, Any]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def generate(
+        self,
+        instruction: str,
+        text: str,
+        examples: list[GenerateDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def parse_force(
+        self,
+        instruction: str,
+        typ: type[T] | dict[str, Any],
+        text: str | None = None,
+        examples: list[ParseForceDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> T | dict[str, Any]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def score(
+        self,
+        instruction: str,
+        text: str,
+        min: int,
+        max: int,
+        examples: list[ScoreDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> int:
+        raise NotImplementedError
+
+
+def default_model() -> GeneralModel:
+    from tau_bench.model_utils.model.openai import OpenAIModel
+
+    return OpenAIModel()
+
+
+def default_quick_model() -> GeneralModel:
+    from tau_bench.model_utils.model.openai import OpenAIModel
+
+    return OpenAIModel(model="gpt-4o-mini")
+
+
+def model_factory(
+    model_id: str,
+    platform: str | Platform,
+    base_url: str | None = None,
+    api_key: str | None = None,
+    temperature: float = 0.0,
+) -> GeneralModel:
+    if isinstance(platform, str):
+        platform = Platform(platform)
+    if platform == Platform.OPENAI:
+        from tau_bench.model_utils.model.openai import OpenAIModel
+
+        return OpenAIModel(model=model_id, api_key=api_key, temperature=temperature)
+    elif platform == Platform.MISTRAL:
+        from tau_bench.model_utils.model.mistral import MistralModel
+
+        return MistralModel(model=model_id, api_key=api_key, temperature=temperature)
+    elif platform == Platform.ANTHROPIC:
+        from tau_bench.model_utils.model.claude import ClaudeModel
+
+        return ClaudeModel(model=model_id, api_key=api_key, temperature=temperature)
+
+    elif platform == Platform.ANYSCALE:
+        from tau_bench.model_utils.model.anyscale import AnyscaleModel
+
+        return AnyscaleModel(model=model_id, api_key=api_key, temperature=temperature)
+    elif platform == Platform.OUTLINES:
+        if base_url is None:
+            raise ValueError("base_url must be provided for custom models")
+        from tau_bench.model_utils.model.outlines_completion import OutlinesCompletionModel
+
+        return OutlinesCompletionModel(model=model_id, base_url=base_url, temperature=temperature)
+    elif platform == Platform.VLLM_CHAT:
+        if base_url is None:
+            raise ValueError("base_url must be provided for custom models")
+        from tau_bench.model_utils.model.vllm_chat import VLLMChatModel
+
+        return VLLMChatModel(
+            model=model_id,
+            base_url=base_url,
+            api_key="not-needed" if api_key is None else api_key,
+            temperature=temperature,
+        )
+    else:
+        if base_url is None:
+            raise ValueError("base_url must be provided for custom models")
+        from tau_bench.model_utils.model.vllm_completion import VLLMCompletionModel
+
+        return VLLMCompletionModel(model=model_id, base_url=base_url, temperature=temperature)

chapter2/prompt-engineering/tau_bench/model_utils/model/mistral.py

@@ -0,0 +1,89 @@
+import os
+
+from tau_bench.model_utils.api.datapoint import Datapoint
+from tau_bench.model_utils.model.chat import ChatModel, Message
+from tau_bench.model_utils.model.completion import approx_cost_for_datapoint, approx_prompt_str
+from tau_bench.model_utils.model.general_model import wrap_temperature
+from tau_bench.model_utils.model.utils import approx_num_tokens
+
+DEFAULT_MISTRAL_MODEL = "mistral-large-latest"
+
+PRICE_PER_INPUT_TOKEN_MAP = {
+    "mistral-largest-latest": 3 / 1000000,
+}
+INPUT_PRICE_PER_TOKEN_FALLBACK = 10 / 1000000
+
+CAPABILITY_SCORE_MAP = {
+    "mistral-largest-latest": 0.9,
+}
+CAPABILITY_SCORE_FALLBACK = 0.3
+
+# TODO: implement
+LATENCY_MS_PER_OUTPUT_TOKEN_MAP = {}
+# TODO: implement
+LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK = 0.0
+
+MAX_CONTEXT_LENGTH_MAP = {
+    "mistral-largest-latest": 128000,
+}
+MAX_CONTEXT_LENGTH_FALLBACK = 128000
+
+
+class MistralModel(ChatModel):
+    def __init__(
+        self, model: str | None = None, api_key: str | None = None, temperature: float = 0.0
+    ) -> None:
+        from mistralai.async_client import MistralAsyncClient
+        from mistralai.client import MistralClient
+
+        if model is None:
+            self.model = DEFAULT_MISTRAL_MODEL
+        else:
+            self.model = model
+
+        api_key = None
+        if api_key is None:
+            api_key = os.getenv("MISTRAL_API_KEY")
+            if api_key is None:
+                raise ValueError("MISTRAL_API_KEY environment variable is not set")
+        self.client = MistralClient(api_key=api_key)
+        self.async_client = MistralAsyncClient(api_key=api_key)
+        self.temperature = temperature
+
+    def generate_message(
+        self,
+        messages: list[Message],
+        force_json: bool,
+        temperature: float | None = None,
+    ) -> Message:
+        if temperature is None:
+            temperature = self.temperature
+        msgs = self.build_generate_message_state(messages)
+        res = self.client.chat(
+            model=self.model,
+            messages=msgs,
+            temperature=wrap_temperature(temperature),
+            response_format={"type": "json_object" if force_json else "text"},
+        )
+        return self.handle_generate_message_response(
+            prompt=msgs, content=res.choices[0].message.content, force_json=force_json
+        )
+
+    def get_approx_cost(self, dp: Datapoint) -> float:
+        cost_per_token = PRICE_PER_INPUT_TOKEN_MAP.get(self.model, INPUT_PRICE_PER_TOKEN_FALLBACK)
+        return approx_cost_for_datapoint(dp=dp, price_per_input_token=cost_per_token)
+
+    def get_latency(self, dp: Datapoint) -> float:
+        latency_per_output_token = LATENCY_MS_PER_OUTPUT_TOKEN_MAP.get(
+            self.model, LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK
+        )
+        return approx_cost_for_datapoint(dp=dp, price_per_input_token=latency_per_output_token)
+
+    def get_capability(self) -> float:
+        return CAPABILITY_SCORE_MAP.get(self.model, CAPABILITY_SCORE_FALLBACK)
+
+    def supports_dp(self, dp: Datapoint) -> bool:
+        prompt = approx_prompt_str(dp)
+        return approx_num_tokens(prompt) <= MAX_CONTEXT_LENGTH_MAP.get(
+            self.model, MAX_CONTEXT_LENGTH_FALLBACK
+        )

chapter2/prompt-engineering/tau_bench/model_utils/model/model.py

@@ -0,0 +1,130 @@
+import abc
+import enum
+from typing import Any, TypeVar
+
+from pydantic import BaseModel
+
+from tau_bench.model_utils.api.datapoint import (
+    BinaryClassifyDatapoint,
+    ClassifyDatapoint,
+    Datapoint,
+    GenerateDatapoint,
+    ParseDatapoint,
+    ParseForceDatapoint,
+    ScoreDatapoint,
+)
+from tau_bench.model_utils.api.types import PartialObj
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class Platform(enum.Enum):
+    OPENAI = "openai"
+    MISTRAL = "mistral"
+    ANTHROPIC = "anthropic"
+    ANYSCALE = "anyscale"
+    OUTLINES = "outlines"
+    VLLM_CHAT = "vllm-chat"
+    VLLM_COMPLETION = "vllm-completion"
+
+
+# @runtime_checkable
+# class Model(Protocol):
+class Model(abc.ABC):
+    @abc.abstractmethod
+    def get_capability(self) -> float:
+        """Return the capability of the model, a float between 0.0 and 1.0."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_approx_cost(self, dp: Datapoint) -> float:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_latency(self, dp: Datapoint) -> float:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def supports_dp(self, dp: Datapoint) -> bool:
+        raise NotImplementedError
+
+
+class ClassifyModel(Model):
+    @abc.abstractmethod
+    def classify(
+        self,
+        instruction: str,
+        text: str,
+        options: list[str],
+        examples: list[ClassifyDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> int:
+        raise NotImplementedError
+
+
+class BinaryClassifyModel(Model):
+    @abc.abstractmethod
+    def binary_classify(
+        self,
+        instruction: str,
+        text: str,
+        examples: list[BinaryClassifyDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> bool:
+        raise NotImplementedError
+
+
+class ParseModel(Model):
+    @abc.abstractmethod
+    def parse(
+        self,
+        text: str,
+        typ: type[T] | dict[str, Any],
+        examples: list[ParseDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> T | PartialObj | dict[str, Any]:
+        raise NotImplementedError
+
+
+class GenerateModel(Model):
+    @abc.abstractmethod
+    def generate(
+        self,
+        instruction: str,
+        text: str,
+        examples: list[GenerateDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> str:
+        raise NotImplementedError
+
+
+class ParseForceModel(Model):
+    @abc.abstractmethod
+    def parse_force(
+        self,
+        instruction: str,
+        typ: type[T] | dict[str, Any],
+        text: str | None = None,
+        examples: list[ParseForceDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> T | dict[str, Any]:
+        raise NotImplementedError
+
+
+class ScoreModel(Model):
+    @abc.abstractmethod
+    def score(
+        self,
+        instruction: str,
+        text: str,
+        min: int,
+        max: int,
+        examples: list[ScoreDatapoint] | None = None,
+        temperature: float | None = None,
+    ) -> int:
+        raise NotImplementedError
+
+
+AnyModel = (
+    BinaryClassifyModel | ClassifyModel | ParseForceModel | GenerateModel | ParseModel | ScoreModel
+)

chapter2/prompt-engineering/tau_bench/model_utils/model/openai.py

@@ -0,0 +1,123 @@
+import os
+
+from tau_bench.model_utils.api.datapoint import Datapoint
+from tau_bench.model_utils.model.chat import ChatModel, Message
+from tau_bench.model_utils.model.completion import approx_cost_for_datapoint, approx_prompt_str
+from tau_bench.model_utils.model.general_model import wrap_temperature
+from tau_bench.model_utils.model.utils import approx_num_tokens
+
+DEFAULT_OPENAI_MODEL = "gpt-4o-2024-08-06"
+API_KEY_ENV_VAR = "OPENAI_API_KEY"
+
+PRICE_PER_INPUT_TOKEN_MAP = {
+    "gpt-4o-2024-08-06": 2.5 / 1000000,
+    "gpt-4o": 5 / 1000000,
+    "gpt-4o-2024-08-06": 2.5 / 1000000,
+    "gpt-4o-2024-05-13": 5 / 1000000,
+    "gpt-4-turbo": 10 / 1000000,
+    "gpt-4-turbo-2024-04-09": 10 / 1000000,
+    "gpt-4": 30 / 1000000,
+    "gpt-4o-mini": 0.15 / 1000000,
+    "gpt-4o-mini-2024-07-18": 0.15 / 1000000,
+    "gpt-3.5-turbo": 0.5 / 1000000,
+    "gpt-3.5-turbo-0125": 0.5 / 1000000,
+    "gpt-3.5-turbo-instruct": 1.5 / 1000000,
+}
+INPUT_PRICE_PER_TOKEN_FALLBACK = 10 / 1000000
+
+CAPABILITY_SCORE_MAP = {
+    "gpt-4o-2024-08-06": 0.8,
+    "gpt-4o": 0.8,
+    "gpt-4o-2024-08-06": 0.8,
+    "gpt-4o-2024-05-13": 0.8,
+    "gpt-4-turbo": 0.9,
+    "gpt-4-turbo-2024-04-09": 0.9,
+    "gpt-4": 0.8,
+    "gpt-4o-mini": 0.5,
+    "gpt-4o-mini-2024-07-18": 0.5,
+    "gpt-3.5-turbo": 0.3,
+    "gpt-3.5-turbo-0125": 0.3,
+}
+CAPABILITY_SCORE_FALLBACK = 0.3
+
+# TODO: implement
+LATENCY_MS_PER_OUTPUT_TOKEN_MAP = {}
+# TODO: implement
+LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK = 0.0
+
+MAX_CONTEXT_LENGTH_MAP = {
+    "gpt-4o-2024-08-06": 128000,
+    "gpt-4o": 128000,
+    "gpt-4o-2024-08-06": 128000,
+    "gpt-4o-2024-05-13": 128000,
+    "gpt-4-turbo": 128000,
+    "gpt-4-turbo-2024-04-09": 128000,
+    "gpt-4": 8192,
+    "gpt-4o-mini": 128000,
+    "gpt-4o-mini-2024-07-18": 128000,
+    "gpt-3.5-turbo": 16385,
+    "gpt-3.5-turbo-0125": 16385,
+}
+MAX_CONTEXT_LENGTH_FALLBACK = 128000
+
+
+class OpenAIModel(ChatModel):
+    def __init__(
+        self,
+        model: str | None = None,
+        api_key: str | None = None,
+        temperature: float = 0.0,
+    ) -> None:
+        from openai import AsyncOpenAI, OpenAI
+
+        if model is None:
+            self.model = DEFAULT_OPENAI_MODEL
+        else:
+            self.model = model
+
+        api_key = None
+        if api_key is None:
+            api_key = os.getenv(API_KEY_ENV_VAR)
+            if api_key is None:
+                raise ValueError(f"{API_KEY_ENV_VAR} environment variable is not set")
+        self.client = OpenAI(api_key=api_key)
+        self.async_client = AsyncOpenAI(api_key=api_key)
+        self.temperature = temperature
+
+    def generate_message(
+        self,
+        messages: list[Message],
+        force_json: bool,
+        temperature: float | None = None,
+    ) -> Message:
+        if temperature is None:
+            temperature = self.temperature
+        msgs = self.build_generate_message_state(messages)
+        res = self.client.chat.completions.create(
+            model=self.model,
+            messages=msgs,
+            temperature=wrap_temperature(temperature),
+            response_format={"type": "json_object" if force_json else "text"},
+        )
+        return self.handle_generate_message_response(
+            prompt=msgs, content=res.choices[0].message.content, force_json=force_json
+        )
+
+    def get_approx_cost(self, dp: Datapoint) -> float:
+        cost_per_token = PRICE_PER_INPUT_TOKEN_MAP.get(self.model, INPUT_PRICE_PER_TOKEN_FALLBACK)
+        return approx_cost_for_datapoint(dp=dp, price_per_input_token=cost_per_token)
+
+    def get_latency(self, dp: Datapoint) -> float:
+        latency_per_output_token = LATENCY_MS_PER_OUTPUT_TOKEN_MAP.get(
+            self.model, LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK
+        )
+        return approx_cost_for_datapoint(dp=dp, price_per_input_token=latency_per_output_token)
+
+    def get_capability(self) -> float:
+        return CAPABILITY_SCORE_MAP.get(self.model, CAPABILITY_SCORE_FALLBACK)
+
+    def supports_dp(self, dp: Datapoint) -> bool:
+        prompt = approx_prompt_str(dp)
+        return approx_num_tokens(prompt) <= MAX_CONTEXT_LENGTH_MAP.get(
+            self.model, MAX_CONTEXT_LENGTH_FALLBACK
+        )

chapter2/prompt-engineering/tau_bench/model_utils/model/outlines_completion.py

@@ -0,0 +1,36 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+from tau_bench.model_utils.api.datapoint import Datapoint
+from tau_bench.model_utils.model.vllm_completion import VLLMCompletionModel
+from tau_bench.model_utils.model.vllm_utils import generate_request
+
+
+class OutlinesCompletionModel(VLLMCompletionModel):
+    def parse_force_from_prompt(
+        self, prompt: str, typ: BaseModel, temperature: float | None = None
+    ) -> dict[str, Any]:
+        if temperature is None:
+            temperature = self.temperature
+        schema = typ.model_json_schema()
+        res = generate_request(
+            url=self.url,
+            prompt=prompt,
+            force_json=True,
+            schema=schema,
+            temperature=temperature,
+        )
+        return self.handle_parse_force_response(prompt=prompt, content=res)
+
+    def get_approx_cost(self, dp: Datapoint) -> float:
+        return super().get_approx_cost(dp)
+
+    def get_latency(self, dp: Datapoint) -> float:
+        return super().get_latency(dp)
+
+    def get_capability(self) -> float:
+        return super().get_capability()
+
+    def supports_dp(self, dp: Datapoint) -> bool:
+        return super().supports_dp(dp)

chapter2/prompt-engineering/tau_bench/model_utils/model/utils.py

@@ -0,0 +1,150 @@
+import enum
+import json
+import re
+from typing import Any, Optional, TypeVar
+
+from pydantic import BaseModel, Field
+
+from tau_bench.model_utils.api.types import PartialObj
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class InputType(enum.Enum):
+    CHAT = "chat"
+    COMPLETION = "completion"
+
+
+def display_choices(choices: list[str]) -> tuple[str, dict[str, int]]:
+    choice_displays = []
+    decode_map = {}
+    for i, choice in enumerate(choices):
+        label = index_to_alpha(i)
+        choice_display = f"{label}. {choice}"
+        choice_displays.append(choice_display)
+        decode_map[label] = i
+    return "\n".join(choice_displays), decode_map
+
+
+def index_to_alpha(index: int) -> str:
+    alpha = ""
+    while index >= 0:
+        alpha = chr(index % 26 + ord("A")) + alpha
+        index = index // 26 - 1
+    return alpha
+
+
+def type_to_json_schema_string(typ: type[T]) -> str:
+    json_schema = typ.model_json_schema()
+    return json.dumps(json_schema, indent=4)
+
+
+def optionalize_type(typ: type[T]) -> type[T]:
+    class OptionalModel(typ):
+        ...
+
+    new_fields = {}
+    for name, field in OptionalModel.model_fields.items():
+        new_fields[name] = Field(default=None, annotation=Optional[field.annotation])
+    OptionalModel.model_fields = new_fields
+    OptionalModel.__name__ = typ.__name__
+    return OptionalModel
+
+
+def json_response_to_obj_or_partial_obj(
+    response: dict[str, Any], typ: type[T] | dict[str, Any]
+) -> T | PartialObj | dict[str, Any]:
+    if isinstance(typ, dict):
+        return response
+    else:
+        required_field_names = [
+            name for name, field in typ.model_fields.items() if field.is_required()
+        ]
+        for name in required_field_names:
+            if name not in response.keys() or response[name] is None:
+                return response
+        return typ.model_validate(response)
+
+
+def clean_top_level_keys(d: dict[str, Any]) -> dict[str, Any]:
+    new_d = {}
+    for k, v in d.items():
+        new_d[k.strip()] = v
+    return new_d
+
+
+def parse_json_or_json_markdown(text: str) -> dict[str, Any]:
+    def parse(s: str) -> dict[str, Any] | None:
+        try:
+            return json.loads(s)
+        except json.decoder.JSONDecodeError:
+            return None
+
+    # pass #1: try to parse as json
+    parsed = parse(text)
+    if parsed is not None:
+        return parsed
+
+    # pass #2: try to parse as json markdown
+    stripped = text.strip()
+    if stripped.startswith("```json"):
+        stripped = stripped[len("```json") :].strip()
+    if stripped.endswith("```"):
+        stripped = stripped[: -len("```")].strip()
+    parsed = parse(stripped)
+    if parsed is not None:
+        return parsed
+
+    # pass #3: try to parse an arbitrary md block
+    pattern = r"```(?:\w+\n)?(.*?)```"
+    match = re.search(pattern, text, re.DOTALL)
+    if match:
+        content = match.group(1).strip()
+        parsed = parse(content)
+        if parsed is not None:
+            return parsed
+
+    # pass #4: try to parse arbitrary sections as json
+    lines = text.split("\n")
+    seen = set()
+    for i in range(len(lines)):
+        for j in range(i + 1, len(lines) + 1):
+            if i < j and (i, j) not in seen:
+                seen.add((i, j))
+                content = "\n".join(lines[i:j])
+                parsed = parse(content)
+                if parsed is not None:
+                    return parsed
+    raise ValueError("Could not parse JSON or JSON markdown")
+
+
+def longest_valid_string(s: str, options: list[str]) -> str | None:
+    longest = 0
+    longest_str = None
+    options_set = set(options)
+    for i in range(len(s)):
+        if s[: i + 1] in options_set and i + 1 > longest:
+            longest = i + 1
+            longest_str = s[: i + 1]
+    return longest_str
+
+
+def try_classify_recover(s: str, decode_map: dict[str, int]) -> str | None:
+    lvs = longest_valid_string(s, list(decode_map.keys()))
+    if lvs is not None and lvs in decode_map:
+        return lvs
+    for k, v in decode_map.items():
+        if s == v:
+            return k
+
+
+def approx_num_tokens(text: str) -> int:
+    return len(text) // 4
+
+
+def add_md_close_tag(prompt: str) -> str:
+    return f"{prompt}\n```"
+
+
+def add_md_tag(prompt: str) -> str:
+    return f"```json\n{prompt}\n```"

chapter2/prompt-engineering/tau_bench/model_utils/model/vllm_chat.py

@@ -0,0 +1,129 @@
+from tau_bench.model_utils.api.datapoint import Datapoint
+from tau_bench.model_utils.model.chat import ChatModel, Message
+from tau_bench.model_utils.model.completion import approx_cost_for_datapoint, approx_prompt_str
+from tau_bench.model_utils.model.general_model import wrap_temperature
+from tau_bench.model_utils.model.utils import approx_num_tokens
+
+PRICE_PER_INPUT_TOKEN_MAP = {
+    "Qwen/Qwen2-0.5B-Instruct": 0.0,
+    "Qwen/Qwen2-1.5B-Instruct": 0.0,
+    "Qwen/Qwen2-7B-Instruct": 0.0,
+    "Qwen/Qwen2-72B-Instruct": 0.0,
+    "meta-llama/Meta-Llama-3.1-8B-Instruct": 0.0,
+    "sierra-research/Meta-Llama-3.1-8B-Instruct": 0.0,
+    "meta-llama/Meta-Llama-3.1-70B-Instruct": 0.0,
+    "mistralai/Mistral-Nemo-Instruct-2407": 0.0,
+}
+INPUT_PRICE_PER_TOKEN_FALLBACK = 0.0
+
+# TODO: refine this
+CAPABILITY_SCORE_MAP = {
+    "Qwen/Qwen2-0.5B-Instruct": 0.05,
+    "Qwen/Qwen2-1.5B-Instruct": 0.07,
+    "Qwen/Qwen2-7B-Instruct": 0.2,
+    "Qwen/Qwen2-72B-Instruct": 0.4,
+    "meta-llama/Meta-Llama-3.1-8B-Instruct": 0.3,
+    "sierra-research/Meta-Llama-3.1-8B-Instruct": 0.3,
+    "meta-llama/Meta-Llama-3.1-70B-Instruct": 0.4,
+    "mistralai/Mistral-Nemo-Instruct-2407": 0.3,
+}
+CAPABILITY_SCORE_FALLBACK = 0.3
+
+# TODO: implement
+LATENCY_MS_PER_OUTPUT_TOKEN_MAP = {}
+# TODO: implement
+LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK = 0.0
+
+MAX_CONTEXT_LENGTH_MAP = {
+    "Qwen/Qwen2-0.5B-Instruct": 32768,
+    "Qwen/Qwen2-1.5B-Instruct": 32768,
+    "Qwen/Qwen2-7B-Instruct": 131072,
+    "Qwen/Qwen2-72B-Instruct": 131072,
+    "meta-llama/Meta-Llama-3.1-8B-Instruct": 128000,
+    "sierra-research/Meta-Llama-3.1-8B-Instruct": 128000,
+    "meta-llama/Meta-Llama-3.1-70B-Instruct": 128000,
+    "mistralai/Mistral-Nemo-Instruct-2407": 128000,
+}
+MAX_CONTEXT_LENGTH_FALLBACK = 128000
+
+
+class VLLMChatModel(ChatModel):
+    def __init__(
+        self,
+        model: str,
+        base_url: str,
+        api_key: str,
+        temperature: float = 0.0,
+        price_per_input_token: float | None = None,
+        capability: float | None = None,
+        latency_ms_per_output_token: float | None = None,
+        max_context_length: int | None = None,
+    ) -> None:
+        from openai import AsyncOpenAI, OpenAI
+
+        self.model = model
+        self.client = OpenAI(
+            base_url=base_url,
+            api_key=api_key,
+        )
+        self.async_client = AsyncOpenAI(
+            base_url=base_url,
+            api_key=api_key,
+        )
+        self.temperature = temperature
+        self.price_per_input_token = (
+            price_per_input_token
+            if price_per_input_token is not None
+            else PRICE_PER_INPUT_TOKEN_MAP.get(model, INPUT_PRICE_PER_TOKEN_FALLBACK)
+        )
+        self.capability = (
+            capability
+            if capability is not None
+            else CAPABILITY_SCORE_MAP.get(model, CAPABILITY_SCORE_FALLBACK)
+        )
+        self.latency_ms_per_output_token = (
+            latency_ms_per_output_token
+            if latency_ms_per_output_token is not None
+            else LATENCY_MS_PER_OUTPUT_TOKEN_MAP.get(model, LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK)
+        )
+        self.max_context_length = (
+            max_context_length
+            if max_context_length is not None
+            else MAX_CONTEXT_LENGTH_MAP.get(model, MAX_CONTEXT_LENGTH_FALLBACK)
+        )
+
+    def get_approx_cost(self, dp: Datapoint) -> float:
+        cost_per_token = self.price_per_input_token
+        return approx_cost_for_datapoint(dp=dp, price_per_input_token=cost_per_token)
+
+    def get_latency(self, dp: Datapoint) -> float:
+        latency_per_output_token = self.latency_ms_per_output_token
+        return approx_cost_for_datapoint(dp=dp, price_per_input_token=latency_per_output_token)
+
+    def get_capability(self) -> float:
+        return CAPABILITY_SCORE_MAP.get(self.model, CAPABILITY_SCORE_FALLBACK)
+
+    def supports_dp(self, dp: Datapoint) -> bool:
+        prompt = approx_prompt_str(dp)
+        return approx_num_tokens(prompt) <= self.max_context_length
+
+    def generate_message(
+        self,
+        messages: list[Message],
+        force_json: bool,
+        temperature: float | None = None,
+    ) -> Message:
+        if temperature is None:
+            temperature = self.temperature
+        msgs = self.build_generate_message_state(messages)
+        res = self.client.chat.completions.create(
+            model=self.model,
+            messages=msgs,
+            temperature=wrap_temperature(temperature=temperature),
+        )
+        return self.handle_generate_message_response(
+            prompt=msgs, content=res.choices[0].message.content, force_json=force_json
+        )
+
+    def force_json_prompt(self, text: str, _: bool = False) -> str:
+        return super().force_json_prompt(text, with_prefix=True)

chapter2/prompt-engineering/tau_bench/model_utils/model/vllm_completion.py

@@ -0,0 +1,121 @@
+import os
+from typing import Any
+
+from pydantic import BaseModel
+
+from tau_bench.model_utils.api.datapoint import Datapoint
+from tau_bench.model_utils.model.completion import (
+    CompletionModel,
+    approx_cost_for_datapoint,
+    approx_prompt_str,
+)
+from tau_bench.model_utils.model.utils import approx_num_tokens
+from tau_bench.model_utils.model.vllm_utils import generate_request
+
+PRICE_PER_INPUT_TOKEN_MAP = {
+    "Qwen/Qwen2-0.5B-Instruct": 0.0,
+    "Qwen/Qwen2-1.5B-Instruct": 0.0,
+    "Qwen/Qwen2-7B-Instruct": 0.0,
+    "Qwen/Qwen2-72B-Instruct": 0.0,
+    "meta-llama/Meta-Llama-3-8B-Instruct": 0.0,
+    "meta-llama/Meta-Llama-3.1-8B-Instruct": 0.0,
+    "meta-llama/Meta-Llama-3-70B-Instruct": 0.0,
+    "mistralai/Mistral-Nemo-Instruct-2407": 0.0,
+}
+INPUT_PRICE_PER_TOKEN_FALLBACK = 0.0
+
+# TODO: refine this
+CAPABILITY_SCORE_MAP = {
+    "Qwen/Qwen2-0.5B-Instruct": 0.05,
+    "Qwen/Qwen2-1.5B-Instruct": 0.07,
+    "Qwen/Qwen2-7B-Instruct": 0.2,
+    "Qwen/Qwen2-72B-Instruct": 0.4,
+    "meta-llama/Meta-Llama-3.1-8B-Instruct": 0.3,
+    "sierra-research/Meta-Llama-3.1-8B-Instruct": 0.3,
+    "meta-llama/Meta-Llama-3.1-70B-Instruct": 0.5,
+    "mistralai/Mistral-Nemo-Instruct-2407": 0.3,
+}
+CAPABILITY_SCORE_FALLBACK = 0.1
+
+# TODO: implement
+LATENCY_MS_PER_OUTPUT_TOKEN_MAP = {}
+# TODO: implement
+LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK = 0.0
+
+MAX_CONTEXT_LENGTH_MAP = {
+    "Qwen/Qwen2-0.5B-Instruct": 32768,
+    "Qwen/Qwen2-1.5B-Instruct": 32768,
+    "Qwen/Qwen2-7B-Instruct": 131072,
+    "Qwen/Qwen2-72B-Instruct": 131072,
+    "meta-llama/Meta-Llama-3.1-8B-Instruct": 128000,
+    "sierra-research/Meta-Llama-3.1-8B-Instruct": 128000,
+    "meta-llama/Meta-Llama-3.1-70B-Instruct": 128000,
+    "mistralai/Mistral-Nemo-Instruct-2407": 128000,
+}
+MAX_CONTEXT_LENGTH_FALLBACK = 128000
+
+
+class VLLMCompletionModel(CompletionModel):
+    def __init__(
+        self,
+        model: str,
+        base_url: str,
+        endpoint: str = "generate",
+        temperature: float = 0.0,
+        price_per_input_token: float | None = None,
+        capability: float | None = None,
+        latency_ms_per_output_token: float | None = None,
+        max_context_length: int | None = None,
+    ) -> None:
+        self.model = model
+        self.base_url = base_url
+        self.url = os.path.join(base_url, endpoint)
+        self.temperature = temperature
+        self.price_per_input_token = (
+            price_per_input_token
+            if price_per_input_token is not None
+            else PRICE_PER_INPUT_TOKEN_MAP.get(model, INPUT_PRICE_PER_TOKEN_FALLBACK)
+        )
+        self.capability = (
+            capability
+            if capability is not None
+            else CAPABILITY_SCORE_MAP.get(model, CAPABILITY_SCORE_FALLBACK)
+        )
+        self.latency_ms_per_output_token = (
+            latency_ms_per_output_token
+            if latency_ms_per_output_token is not None
+            else LATENCY_MS_PER_OUTPUT_TOKEN_MAP.get(model, LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK)
+        )
+        self.max_context_length = (
+            max_context_length
+            if max_context_length is not None
+            else MAX_CONTEXT_LENGTH_MAP.get(model, MAX_CONTEXT_LENGTH_FALLBACK)
+        )
+
+    def generate_from_prompt(self, prompt: str, temperature: float = 0.0) -> str:
+        return generate_request(url=self.url, prompt=prompt, temperature=temperature)
+
+    def parse_force_from_prompt(
+        self, prompt: str, typ: BaseModel | dict[str, Any], temperature: float | None = None
+    ) -> dict[str, Any]:
+        if temperature is None:
+            temperature = self.temperature
+        res = generate_request(
+            url=self.url, prompt=prompt, force_json=True, temperature=temperature
+        )
+        return self.handle_parse_force_response(prompt=prompt, content=res)
+
+    def get_approx_cost(self, dp: Datapoint) -> float:
+        cost_per_token = self.price_per_input_token
+        return approx_cost_for_datapoint(dp=dp, price_per_input_token=cost_per_token)
+
+    def get_latency(self, dp: Datapoint) -> float:
+        latency_per_output_token = self.latency_ms_per_output_token
+        return approx_cost_for_datapoint(dp=dp, price_per_input_token=latency_per_output_token)
+
+    def get_capability(self) -> float:
+        return CAPABILITY_SCORE_MAP.get(self.model, CAPABILITY_SCORE_FALLBACK)
+
+    def supports_dp(self, dp: Datapoint) -> bool:
+        prompt = approx_prompt_str(dp)
+        return approx_num_tokens(prompt) <= self.max_context_length

chapter2/prompt-engineering/tau_bench/model_utils/model/vllm_utils.py

@@ -0,0 +1,36 @@
+from typing import Any
+
+import requests
+
+from tau_bench.model_utils.model.general_model import wrap_temperature
+
+
+def generate_request(
+    url: str,
+    prompt: str,
+    temperature: float = 0.0,
+    force_json: bool = False,
+    **req_body_kwargs: Any,
+) -> str:
+    args = {
+        "prompt": prompt,
+        "temperature": wrap_temperature(temperature),
+        "max_tokens": 4096,
+        **req_body_kwargs,
+    }
+    if force_json:
+        # the prompt will have a suffix of '```json\n' to indicate that the response should be a JSON object
+        args["stop"] = ["```"]
+    res = requests.post(
+        url,
+        json=args,
+    )
+    res.raise_for_status()
+    json_res = res.json()
+    if "text" not in json_res:
+        raise ValueError(f"Unexpected response: {json_res}")
+    elif len(json_res["text"]) == 0:
+        raise ValueError(f"Empty response: {json_res}")
+    text = json_res["text"][0]
+    assert isinstance(text, str)
+    return text.removeprefix(prompt)