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

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

@@ -0,0 +1,50 @@
+from tau_bench.model_utils.api.api import API as API
+from tau_bench.model_utils.api.api import default_api_from_args as default_api_from_args
+from tau_bench.model_utils.api.api import BinaryClassifyDatapoint as BinaryClassifyDatapoint
+from tau_bench.model_utils.api.api import ClassifyDatapoint as ClassifyDatapoint
+from tau_bench.model_utils.api.api import GenerateDatapoint as GenerateDatapoint
+from tau_bench.model_utils.api.api import ParseDatapoint as ParseDatapoint
+from tau_bench.model_utils.api.api import ParseForceDatapoint as ParseForceDatapoint
+from tau_bench.model_utils.api.api import ScoreDatapoint as ScoreDatapoint
+from tau_bench.model_utils.api.api import default_api as default_api
+from tau_bench.model_utils.api.api import default_quick_api as default_quick_api
+from tau_bench.model_utils.api.datapoint import Datapoint as Datapoint
+from tau_bench.model_utils.api.datapoint import EvaluationResult as EvaluationResult
+from tau_bench.model_utils.api.datapoint import datapoint_factory as datapoint_factory
+from tau_bench.model_utils.api.datapoint import load_from_disk as load_from_disk
+from tau_bench.model_utils.api.exception import APIError as APIError
+from tau_bench.model_utils.api.sample import (
+    EnsembleSamplingStrategy as EnsembleSamplingStrategy,
+)
+from tau_bench.model_utils.api.sample import (
+    MajoritySamplingStrategy as MajoritySamplingStrategy,
+)
+from tau_bench.model_utils.api.sample import (
+    RedundantSamplingStrategy as RedundantSamplingStrategy,
+)
+from tau_bench.model_utils.api.sample import RetrySamplingStrategy as RetrySamplingStrategy
+from tau_bench.model_utils.api.sample import SamplingStrategy as SamplingStrategy
+from tau_bench.model_utils.api.sample import SingleSamplingStrategy as SingleSamplingStrategy
+from tau_bench.model_utils.api.sample import (
+    UnanimousSamplingStrategy as UnanimousSamplingStrategy,
+)
+from tau_bench.model_utils.api.sample import (
+    get_default_sampling_strategy as get_default_sampling_strategy,
+)
+from tau_bench.model_utils.api.sample import (
+    set_default_sampling_strategy as set_default_sampling_strategy,
+)
+from tau_bench.model_utils.model.chat import PromptSuffixStrategy as PromptSuffixStrategy
+from tau_bench.model_utils.model.exception import ModelError as ModelError
+from tau_bench.model_utils.model.general_model import GeneralModel as GeneralModel
+from tau_bench.model_utils.model.general_model import default_model as default_model
+from tau_bench.model_utils.model.general_model import model_factory as model_factory
+from tau_bench.model_utils.model.model import BinaryClassifyModel as BinaryClassifyModel
+from tau_bench.model_utils.model.model import ClassifyModel as ClassifyModel
+from tau_bench.model_utils.model.model import GenerateModel as GenerateModel
+from tau_bench.model_utils.model.model import ParseForceModel as ParseForceModel
+from tau_bench.model_utils.model.model import ParseModel as ParseModel
+from tau_bench.model_utils.model.model import Platform as Platform
+from tau_bench.model_utils.model.model import ScoreModel as ScoreModel
+from tau_bench.model_utils.model.openai import OpenAIModel as OpenAIModel
+from tau_bench.model_utils.model.utils import InputType as InputType

chapter2/prompt-engineering/tau_bench/model_utils/api/_model_methods.py

@@ -0,0 +1,8 @@
+MODEL_METHODS = [
+    "classify",
+    "binary_classify",
+    "parse",
+    "generate",
+    "parse_force",
+    "score",
+]

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

@@ -0,0 +1,432 @@
+from __future__ import annotations
+
+import argparse
+from typing import Any, TypeVar
+
+from pydantic import BaseModel
+
+from tau_bench.model_utils.api._model_methods import MODEL_METHODS
+from tau_bench.model_utils.api.cache import cache_call_w_dedup
+from tau_bench.model_utils.api.datapoint import (
+    BinaryClassifyDatapoint,
+    ClassifyDatapoint,
+    Datapoint,
+    GenerateDatapoint,
+    ParseDatapoint,
+    ParseForceDatapoint,
+    ScoreDatapoint,
+)
+from tau_bench.model_utils.api.logging import log_call
+from tau_bench.model_utils.api.router import RequestRouter, default_request_router
+from tau_bench.model_utils.api.sample import (
+    EnsembleSamplingStrategy,
+    MajoritySamplingStrategy,
+    SamplingStrategy,
+    get_default_sampling_strategy,
+)
+from tau_bench.model_utils.api.types import PartialObj
+from tau_bench.model_utils.model.general_model import GeneralModel
+from tau_bench.model_utils.model.model import (
+    AnyModel,
+    BinaryClassifyModel,
+    ClassifyModel,
+    GenerateModel,
+    ParseForceModel,
+    ParseModel,
+    ScoreModel,
+)
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class API(object):
+    wrappers_for_main_methods = [log_call, cache_call_w_dedup]
+
+    def __init__(
+        self,
+        parse_models: list[ParseModel],
+        generate_models: list[GenerateModel],
+        parse_force_models: list[ParseForceModel],
+        score_models: list[ScoreModel],
+        classify_models: list[ClassifyModel],
+        binary_classify_models: list[BinaryClassifyModel] | None = None,
+        sampling_strategy: SamplingStrategy | None = None,
+        request_router: RequestRouter | None = None,
+        log_file: str | None = None,
+    ) -> None:
+        if sampling_strategy is None:
+            sampling_strategy = get_default_sampling_strategy()
+        if request_router is None:
+            request_router = default_request_router()
+        self.sampling_strategy = sampling_strategy
+        self.request_router = request_router
+        self._log_file = log_file
+        self.binary_classify_models = binary_classify_models
+        self.classify_models = classify_models
+        self.parse_models = parse_models
+        self.generate_models = generate_models
+        self.parse_force_models = parse_force_models
+        self.score_models = score_models
+
+        self.__init_subclass__()
+
+        self.__init_subclass__()
+
+    def __init_subclass__(cls):
+        for method_name in MODEL_METHODS:
+            if hasattr(cls, method_name):
+                method = getattr(cls, method_name)
+                for wrapper in cls.wrappers_for_main_methods:
+                    method = wrapper(method)
+                setattr(cls, method_name, method)
+
+    @classmethod
+    def from_general_model(
+        cls,
+        model: GeneralModel,
+        sampling_strategy: SamplingStrategy | None = None,
+        request_router: RequestRouter | None = None,
+        log_file: str | None = None,
+    ) -> "API":
+        return cls(
+            binary_classify_models=[model],
+            classify_models=[model],
+            parse_models=[model],
+            generate_models=[model],
+            parse_force_models=[model],
+            score_models=[model],
+            log_file=log_file,
+            sampling_strategy=sampling_strategy,
+            request_router=request_router,
+        )
+
+    @classmethod
+    def from_general_models(
+        cls,
+        models: list[GeneralModel],
+        sampling_strategy: SamplingStrategy | None = None,
+        request_router: RequestRouter | None = None,
+        log_file: str | None = None,
+    ) -> "API":
+        if len(models) == 0:
+            raise ValueError("Must provide at least one model")
+        return cls(
+            binary_classify_models=models,
+            classify_models=models,
+            parse_models=models,
+            generate_models=models,
+            parse_force_models=models,
+            score_models=models,
+            log_file=log_file,
+            sampling_strategy=sampling_strategy,
+            request_router=request_router,
+        )
+
+    def set_default_binary_classify_models(self, models: list[BinaryClassifyModel]) -> None:
+        if len(models) == 0:
+            raise ValueError("Must provide at least one model")
+        self.binary_classify_models = models
+
+    def set_default_classify_models(self, models: list[BinaryClassifyModel]) -> None:
+        if len(models) == 0:
+            raise ValueError("Must provide at least one model")
+        self.classify_models = models
+
+    def set_default_parse_models(self, models: list[ParseModel]) -> None:
+        if len(models) == 0:
+            raise ValueError("Must provide at least one model")
+        self.parse_models = models
+
+    def set_default_generate_models(self, models: list[GenerateModel]) -> None:
+        if len(models) == 0:
+            raise ValueError("Must provide at least one model")
+        self.generate_models = models
+
+    def set_default_parse_force_models(self, models: list[ParseForceModel]) -> None:
+        if len(models) == 0:
+            raise ValueError("Must provide at least one model")
+        self.parse_force_models = models
+
+    def set_default_score_models(self, models: list[ScoreModel]) -> None:
+        if len(models) == 0:
+            raise ValueError("Must provide at least one model")
+        self.score_models = models
+
+    def set_default_sampling_strategy(self, sampling_strategy: SamplingStrategy) -> None:
+        self.sampling_strategy = sampling_strategy
+
+    def set_default_request_router(self, request_router: RequestRouter) -> None:
+        self.request_router = request_router
+
+    def _run_with_sampling_strategy(
+        self,
+        models: list[AnyModel],
+        datapoint: Datapoint,
+        sampling_strategy: SamplingStrategy,
+    ) -> T:
+        assert len(models) > 0
+
+        def _run_datapoint(model: AnyModel, temp: float | None = None) -> T:
+            if isinstance(datapoint, ClassifyDatapoint):
+                return model.classify(
+                    instruction=datapoint.instruction,
+                    text=datapoint.text,
+                    options=datapoint.options,
+                    examples=datapoint.examples,
+                    temperature=temp,
+                )
+            elif isinstance(datapoint, BinaryClassifyDatapoint):
+                return model.binary_classify(
+                    instruction=datapoint.instruction,
+                    text=datapoint.text,
+                    examples=datapoint.examples,
+                    temperature=temp,
+                )
+            elif isinstance(datapoint, ParseForceDatapoint):
+                return model.parse_force(
+                    instruction=datapoint.instruction,
+                    typ=datapoint.typ,
+                    text=datapoint.text,
+                    examples=datapoint.examples,
+                    temperature=temp,
+                )
+            elif isinstance(datapoint, GenerateDatapoint):
+                return model.generate(
+                    instruction=datapoint.instruction,
+                    text=datapoint.text,
+                    examples=datapoint.examples,
+                    temperature=temp,
+                )
+            elif isinstance(datapoint, ParseDatapoint):
+                return model.parse(
+                    text=datapoint.text,
+                    typ=datapoint.typ,
+                    examples=datapoint.examples,
+                    temperature=temp,
+                )
+            elif isinstance(datapoint, ScoreDatapoint):
+                return model.score(
+                    instruction=datapoint.instruction,
+                    text=datapoint.text,
+                    min=datapoint.min,
+                    max=datapoint.max,
+                    examples=datapoint.examples,
+                    temperature=temp,
+                )
+            else:
+                raise ValueError(f"Unknown datapoint type: {type(datapoint)}")
+
+        if isinstance(sampling_strategy, EnsembleSamplingStrategy):
+            return sampling_strategy.execute(
+                [lambda x=model: _run_datapoint(x, 0.0) for model in models]
+            )
+        return sampling_strategy.execute(
+            lambda: _run_datapoint(
+                models[0], 0.2 if isinstance(sampling_strategy, MajoritySamplingStrategy) else None
+            )
+        )
+
+    def _api_call(
+        self, models: list[AnyModel], datapoint: Datapoint, sampling_strategy: SamplingStrategy
+    ) -> T:
+        if isinstance(sampling_strategy, EnsembleSamplingStrategy):
+            return self._run_with_sampling_strategy(models, datapoint, sampling_strategy)
+        model = self.request_router.route(dp=datapoint, available_models=models)
+        return self._run_with_sampling_strategy(
+            models=[model], datapoint=datapoint, sampling_strategy=sampling_strategy
+        )
+
+    def classify(
+        self,
+        instruction: str,
+        text: str,
+        options: list[str],
+        examples: list[ClassifyDatapoint] | None = None,
+        sampling_strategy: SamplingStrategy | None = None,
+        request_router: RequestRouter | None = None,
+        models: list[ClassifyModel] | None = None,
+    ) -> int:
+        if models is None:
+            models = self.classify_models
+        if sampling_strategy is None:
+            sampling_strategy = self.sampling_strategy
+        if request_router is None:
+            request_router = self.request_router
+
+        return self._api_call(
+            models=models,
+            datapoint=ClassifyDatapoint(
+                instruction=instruction, text=text, options=options, examples=examples
+            ),
+            sampling_strategy=sampling_strategy,
+        )
+
+    def binary_classify(
+        self,
+        instruction: str,
+        text: str,
+        examples: list[BinaryClassifyDatapoint] | None = None,
+        sampling_strategy: SamplingStrategy | None = None,
+        request_router: RequestRouter | None = None,
+        models: list[BinaryClassifyModel] | None = None,
+    ) -> bool:
+        if models is None:
+            models = (
+                self.binary_classify_models
+                if self.binary_classify_models is not None
+                else self.classify_models
+            )
+        if sampling_strategy is None:
+            sampling_strategy = self.sampling_strategy
+        if request_router is None:
+            request_router = self.request_router
+
+        return self._api_call(
+            models=models,
+            datapoint=BinaryClassifyDatapoint(
+                instruction=instruction, text=text, examples=examples
+            ),
+            sampling_strategy=sampling_strategy,
+        )
+
+    def parse(
+        self,
+        text: str,
+        typ: type[T] | dict[str, Any],
+        examples: list[ParseDatapoint] | None = None,
+        sampling_strategy: SamplingStrategy | None = None,
+        request_router: RequestRouter | None = None,
+        models: list[ParseModel] | None = None,
+    ) -> T | PartialObj | dict[str, Any]:
+        if models is None:
+            models = self.parse_models
+        if sampling_strategy is None:
+            sampling_strategy = self.sampling_strategy
+        if request_router is None:
+            request_router = self.request_router
+
+        return self._api_call(
+            models=models,
+            datapoint=ParseDatapoint(text=text, typ=typ, examples=examples),
+            sampling_strategy=sampling_strategy,
+        )
+
+    def generate(
+        self,
+        instruction: str,
+        text: str,
+        examples: list[GenerateDatapoint] | None = None,
+        sampling_strategy: SamplingStrategy | None = None,
+        request_router: RequestRouter | None = None,
+        models: list[GenerateModel] | None = None,
+    ) -> str:
+        if models is None:
+            models = self.generate_models
+        if sampling_strategy is None:
+            sampling_strategy = self.sampling_strategy
+        if request_router is None:
+            request_router = self.request_router
+
+        return self._api_call(
+            models=models,
+            datapoint=GenerateDatapoint(instruction=instruction, text=text, examples=examples),
+            sampling_strategy=sampling_strategy,
+        )
+
+    def parse_force(
+        self,
+        instruction: str,
+        typ: type[T] | dict[str, Any],
+        text: str | None = None,
+        examples: list[ParseForceDatapoint] | None = None,
+        sampling_strategy: SamplingStrategy | None = None,
+        request_router: RequestRouter | None = None,
+        models: list[ParseForceModel] | None = None,
+    ) -> T | dict[str, Any]:
+        if models is None:
+            models = self.parse_force_models
+        if sampling_strategy is None:
+            sampling_strategy = self.sampling_strategy
+        if request_router is None:
+            request_router = self.request_router
+
+        return self._api_call(
+            models=models,
+            datapoint=ParseForceDatapoint(
+                instruction=instruction, typ=typ, text=text, examples=examples
+            ),
+            sampling_strategy=sampling_strategy,
+        )
+
+    def score(
+        self,
+        instruction: str,
+        text: str,
+        min: int,
+        max: int,
+        examples: list[ScoreDatapoint] | None = None,
+        sampling_strategy: SamplingStrategy | None = None,
+        request_router: RequestRouter | None = None,
+        models: list[ScoreModel] | None = None,
+    ) -> int:
+        if models is None:
+            models = self.score_models
+        if sampling_strategy is None:
+            sampling_strategy = self.sampling_strategy
+        if request_router is None:
+            request_router = self.request_router
+
+        return self._api_call(
+            models=models,
+            datapoint=ScoreDatapoint(
+                instruction=instruction, text=text, min=min, max=max, examples=examples
+            ),
+            sampling_strategy=sampling_strategy,
+        )
+
+
+def default_api(
+    log_file: str | None = None,
+    sampling_strategy: SamplingStrategy | None = None,
+    request_router: RequestRouter | None = None,
+) -> API:
+    from tau_bench.model_utils.model.general_model import default_model
+
+    model = default_model()
+    return API(
+        binary_classify_models=[model],
+        classify_models=[model],
+        parse_models=[model],
+        generate_models=[model],
+        parse_force_models=[model],
+        score_models=[model],
+        sampling_strategy=sampling_strategy,
+        request_router=request_router,
+        log_file=log_file,
+    )
+
+def default_api_from_args(args: argparse.Namespace) -> API:
+    from tau_bench.model_utils.model.general_model import model_factory
+    model = model_factory(model_id=args.model, platform=args.platform, base_url=args.base_url)
+    return API.from_general_model(model=model)
+
+
+def default_quick_api(
+    log_file: str | None = None,
+    sampling_strategy: SamplingStrategy | None = None,
+    request_router: RequestRouter | None = None,
+) -> API:
+    from tau_bench.model_utils.model.general_model import default_quick_model
+
+    model = default_quick_model()
+    return API(
+        binary_classify_models=[model],
+        classify_models=[model],
+        parse_models=[model],
+        generate_models=[model],
+        parse_force_models=[model],
+        score_models=[model],
+        sampling_strategy=sampling_strategy,
+        request_router=request_router,
+        log_file=log_file,
+    )

chapter2/prompt-engineering/tau_bench/model_utils/api/cache.py

@@ -0,0 +1,115 @@
+import functools
+import inspect
+import threading
+from collections import defaultdict
+from multiprocessing import Lock
+from typing import Any, Callable, TypeVar
+
+from pydantic import BaseModel
+
+T = TypeVar("T")
+
+
+class _CallableIdentity:
+    __slots__ = ("func",)
+
+    def __init__(self, func: Callable[..., Any]):
+        self.func = func
+
+    def __hash__(self) -> int:
+        return id(self.func)
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, _CallableIdentity) and self.func is other.func
+
+
+CacheKey = tuple[_CallableIdentity, Any]
+
+USE_CACHE = True
+_USE_CACHE_LOCK = Lock()
+cache: dict[CacheKey, tuple[T, threading.Event]] = {}
+lock = threading.Lock()
+conditions = defaultdict(threading.Condition)
+
+
+def disable_cache():
+    global USE_CACHE
+    with _USE_CACHE_LOCK:
+        USE_CACHE = False
+
+
+def enable_cache():
+    global USE_CACHE
+    with _USE_CACHE_LOCK:
+        USE_CACHE = True
+
+
+def hash_item(item: Any) -> Any:
+    if isinstance(item, dict):
+        return (
+            "dict",
+            frozenset(
+                (hash_item(key), hash_item(value)) for key, value in item.items()
+            ),
+        )
+    elif isinstance(item, list):
+        return ("list", tuple(hash_item(x) for x in item))
+    elif isinstance(item, set):
+        return (
+            "set",
+            frozenset(hash_item(x) for x in item),
+        )
+    elif isinstance(item, tuple):
+        return ("tuple", tuple(hash_item(x) for x in item))
+    elif isinstance(item, BaseModel):
+        values = item.model_dump() if hasattr(item, "model_dump") else item.dict()
+        return (
+            "model",
+            type(item).__module__,
+            type(item).__qualname__,
+            hash_item(values),
+        )
+    return item
+
+
+def hash_func_call(
+    func: Callable[..., Any], args: tuple[Any], kwargs: dict[str, Any]
+) -> CacheKey:
+    bound_args = inspect.signature(func).bind(*args, **kwargs)
+    bound_args.apply_defaults()
+    standardized_args = sorted(bound_args.arguments.items())
+    return _CallableIdentity(func), hash_item(standardized_args)
+
+
+def cache_call_w_dedup(func: Callable[..., T]) -> Callable[..., T]:
+    @functools.wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> T:
+        if not USE_CACHE:
+            return func(*args, **kwargs)
+        key = hash_func_call(func=func, args=args, kwargs=kwargs)
+        if key in cache:
+            result, event = cache[key]
+            if event.is_set():
+                return result
+        else:
+            with lock:
+                cache[key] = (None, threading.Event())
+
+        condition = conditions[key]
+        with condition:
+            if cache[key][1].is_set():
+                return cache[key][0]
+            if not cache[key][0]:
+                try:
+                    result = func(*args, **kwargs)
+                    with lock:
+                        cache[key] = (result, threading.Event())
+                        cache[key][1].set()
+                except Exception as e:
+                    with lock:
+                        cache[key] = (e, threading.Event())
+                        cache[key][1].set()
+                    raise e
+            return cache[key][0]
+
+    return wrapper

chapter2/prompt-engineering/tau_bench/model_utils/api/datapoint.py

@@ -0,0 +1,299 @@
+from __future__ import annotations
+
+import abc
+import json
+from typing import Any, Callable, TypeVar
+
+from pydantic import BaseModel
+
+import tau_bench.model_utils
+from tau_bench.model_utils.api._model_methods import MODEL_METHODS
+from tau_bench.model_utils.api.exception import APIError
+from tau_bench.model_utils.api.types import PartialObj
+from tau_bench.model_utils.model.exception import ModelError
+
+T = TypeVar("T", bound=BaseModel)
+
+
+def _is_trace(obj: dict[str, Any]) -> bool:
+    return (
+        "method_name" in obj
+        and obj["method_name"] in MODEL_METHODS
+        and "kwargs" in obj
+        and "response" in obj
+        and isinstance(obj["kwargs"], dict)
+    )
+
+
+def dict_equal(d1: dict, d2: dict) -> bool:
+    d1_keys_sorted = sorted(d1.keys())
+    d2_keys_sorted = sorted(d2.keys())
+    if d1_keys_sorted != d2_keys_sorted:
+        return False
+    for k in d1_keys_sorted:
+        if isinstance(d1[k], dict) and isinstance(d2[k], dict):
+            if not dict_equal(d1[k], d2[k]):
+                return False
+        elif isinstance(d1[k], list) and isinstance(d2[k], list):
+            if not list_equal(d1[k], d2[k]):
+                return False
+        elif isinstance(d1[k], set) and isinstance(d2[k], set):
+            if d1[k] != d2[k]:
+                return False
+        elif isinstance(d1[k], str) and isinstance(d2[k], str):
+            if not str_equal(d1[k], d2[k]):
+                return False
+        elif d1[k] != d2[k]:
+            return False
+    return True
+
+
+def list_equal(l1: list, l2: list) -> bool:
+    if len(l1) != len(l2):
+        return False
+    for i1, i2 in zip(l1, l2):
+        if isinstance(i1, dict) and isinstance(i2, dict):
+            if not dict_equal(i1, i2):
+                return False
+        elif isinstance(i1, list) and isinstance(i2, list):
+            if not list_equal(i1, i2):
+                return False
+        elif isinstance(i1, set) and isinstance(i2, set):
+            if i1 != i2:
+                return False
+        elif isinstance(i1, str) and isinstance(i2, str):
+            if not str_equal(i1, i2):
+                return False
+        elif i1 != i2:
+            return False
+    return True
+
+
+def set_equal(s1: set, s2: set) -> bool:
+    if len(s1) != len(s2):
+        return False
+    for i1, i2 in zip(s1, s2):
+        if isinstance(i1, dict) and isinstance(i2, dict):
+            if not dict_equal(i1, i2):
+                return False
+        elif isinstance(i1, list) and isinstance(i2, list):
+            if not list_equal(i1, i2):
+                return False
+        elif isinstance(i1, set) and isinstance(i2, set):
+            if i1 != i2:
+                return False
+        elif isinstance(i1, str) and isinstance(i2, str):
+            if not str_equal(i1, i2):
+                return False
+        elif i1 != i2:
+            return False
+    return True
+
+
+def str_equal(s1: str, s2: str) -> bool:
+    def remove_special_chars(s: str) -> str:
+        return "".join(filter(str.isalnum, s))
+
+    def strip_and_lower(s: str) -> str:
+        return s.lower().strip()
+
+    return strip_and_lower(remove_special_chars(s1)) == strip_and_lower(remove_special_chars(s2))
+
+
+class EvaluationResult(BaseModel):
+    is_error: bool
+    is_correct: bool
+    datapoint: dict[str, Any] | None
+    response: Any | None
+    error: str | None
+
+
+class Datapoint(BaseModel, abc.ABC):
+    @classmethod
+    def from_trace(cls, d: dict[str, Any]) -> "Datapoint":
+        if not _is_trace(d):
+            raise ValueError(f"This is not a trace: {d}")
+        response = d["response"]
+        kwargs = d["kwargs"]
+        return cls(response=response, **kwargs)
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "Datapoint":
+        if _is_trace(d):
+            return cls.from_trace(d)
+        return cls(**d)
+
+    @abc.abstractmethod
+    def evaluate(self, api: tau_bench.model_utils.API) -> EvaluationResult:
+        raise NotImplementedError
+
+
+class ClassifyDatapoint(Datapoint):
+    instruction: str
+    text: str
+    options: list[str]
+    response: int | None = None
+    examples: list["ClassifyDatapoint"] | None = None
+
+    def evaluate(self, api: tau_bench.model_utils.API) -> EvaluationResult:
+        return run_and_catch_api_error(
+            lambda: api.classify(
+                instruction=self.instruction,
+                text=self.text,
+                options=self.options,
+                examples=self.examples,
+            ),
+            self.response,
+            self.model_dump(),
+        )
+
+
+class BinaryClassifyDatapoint(Datapoint):
+    instruction: str
+    text: str
+    response: bool | None = None
+    examples: list["BinaryClassifyDatapoint"] | None = None
+
+    def evaluate(self, api: tau_bench.model_utils.API) -> EvaluationResult:
+        return run_and_catch_api_error(
+            lambda: api.binary_classify(
+                instruction=self.instruction, text=self.text, examples=self.examples
+            ),
+            self.response,
+            self.model_dump(),
+        )
+
+
+class ScoreDatapoint(Datapoint):
+    instruction: str
+    text: str
+    min: int
+    max: int
+    response: int | None = None
+    examples: list["ScoreDatapoint"] | None = None
+
+    def evaluate(self, api: tau_bench.model_utils.API) -> EvaluationResult:
+        raise NotImplementedError
+
+
+class ParseDatapoint(Datapoint):
+    text: str
+    typ: type[T] | dict[str, Any]
+    response: dict[str, Any] | T | PartialObj | None = None
+    examples: list["ParseDatapoint"] | None = None
+
+    def evaluate(self, api: tau_bench.model_utils.API) -> EvaluationResult:
+        return run_and_catch_api_error(
+            lambda: api.parse(text=self.text, typ=self.typ),
+            self.response,
+            self.model_dump(),
+        )
+
+
+class GenerateDatapoint(Datapoint):
+    instruction: str
+    text: str
+    response: str | None = None
+    examples: list["GenerateDatapoint"] | None = None
+
+    def evaluate(self, api: tau_bench.model_utils.API) -> tau_bench.model_utils.EvaluationResult:
+        raise NotImplementedError
+
+
+class ParseForceDatapoint(Datapoint):
+    instruction: str
+    typ: type[T] | dict[str, Any]
+    text: str | None = None
+    response: dict[str, Any] | T | None = None
+    examples: list["ParseForceDatapoint"] | None = None
+
+    def evaluate(self, api: tau_bench.model_utils.API) -> EvaluationResult:
+        return run_and_catch_api_error(
+            lambda: api.parse_force(
+                instruction=self.instruction,
+                text=self.text,
+                typ=self.typ,
+                examples=self.examples,
+            ),
+            self.response,
+            self.model_dump(),
+        )
+
+
+def datapoint_factory(d: dict[str, Any]) -> Datapoint:
+    if _is_trace(d):
+        method_name = d["method_name"]
+        kwargs = d["kwargs"]
+        data = {"response": d["response"], **kwargs}
+        if method_name == "classify":
+            return ClassifyDatapoint(**data)
+        elif method_name == "binary_classify":
+            return BinaryClassifyDatapoint(**data)
+        elif method_name == "parse":
+            return ParseDatapoint(**data)
+        elif method_name == "parse_force":
+            return ParseForceDatapoint(**data)
+        elif method_name == "generate":
+            return GenerateDatapoint(**data)
+        elif method_name == "score":
+            return ScoreDatapoint(**data)
+        else:
+            raise ValueError(f"Unknown method name: {method_name}")
+    else:
+        if all(k in d for k in ["instruction", "text", "options"]) and isinstance(
+            d["response"], int
+        ):
+            return ClassifyDatapoint(**d)
+        elif all(k in d for k in ["instruction", "text"]) and isinstance(d["response"], bool):
+            return BinaryClassifyDatapoint(**d)
+        elif all(k in d for k in ["instruction", "text", "min", "max"]) and isinstance(
+            d["response"], int
+        ):
+            return ScoreDatapoint(**d)
+        elif all(k in d for k in ["instruction", "text", "typ"]) and isinstance(
+            d["response"], dict
+        ):
+            return ParseForceDatapoint(**d)
+        elif all(k in d for k in ["text", "typ"]) and isinstance(d["response"], dict):
+            return ParseDatapoint(**d)
+        elif all(k in d for k in ["instruction", "text"]) and isinstance(d["response"], str):
+            return GenerateDatapoint(**d)
+        else:
+            raise ValueError(f"Unknown datapoint: {d}")
+
+
+def run_and_catch_api_error(
+    callable: Callable[..., Any], response: Any, datapoint: dict[str, Any]
+) -> EvaluationResult:
+    try:
+        res = callable()
+        if isinstance(response, dict):
+            is_correct = dict_equal(res, response)
+        else:
+            is_correct = res == response
+        return EvaluationResult(
+            is_error=False,
+            is_correct=is_correct,
+            response=res,
+            error=None,
+            datapoint=datapoint,
+        )
+    except (APIError, ModelError) as e:
+        return EvaluationResult(
+            is_error=True,
+            is_correct=False,
+            response=None,
+            error=str(e),
+            datapoint=datapoint,
+        )
+
+
+def load_from_disk(path: str) -> list[Datapoint]:
+    with open(path, "r") as f:
+        if path.endswith(".jsonl"):
+            data = [json.loads(line) for line in f]
+        elif path.endswith(".json"):
+            data = json.load(f)
+        else:
+            raise ValueError(f"Unknown file format: {path}")
+    return [datapoint_factory(d) for d in data]

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

@@ -0,0 +1,69 @@
+import json
+import os
+import time
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Callable, TypeVar
+
+from tau_bench.model_utils.model.exception import ModelError, Result
+
+T = TypeVar("T")
+
+_REPORT_DIR = os.path.expanduser("~/.llm-primitives/log")
+
+
+def set_report_dir(path: str) -> None:
+    global _REPORT_DIR
+    _REPORT_DIR = path
+
+
+def get_report_dir() -> str:
+    return _REPORT_DIR
+
+
+def log_report_to_disk(report: dict[str, Any], path: str) -> None:
+    with open(path, "w") as f:
+        json.dump(report, f, indent=4)
+
+
+def generate_report_location() -> str:
+    if not os.path.exists(_REPORT_DIR):
+        os.makedirs(_REPORT_DIR)
+    return os.path.join(_REPORT_DIR, f"report-{time.time_ns()}.json")
+
+
+class APIError(Exception):
+    def __init__(self, short_message: str, report: dict[str, Any] | None = None) -> None:
+        self.report_path = generate_report_location()
+        self.short_message = short_message
+        self.report = report
+        if self.report is not None:
+            log_report_to_disk(
+                report={"error_type": "APIError", "report": report}, path=self.report_path
+            )
+        super().__init__(f"{short_message}\n\nSee the full report at {self.report_path}")
+
+
+def execute_and_filter_model_errors(
+    funcs: list[Callable[[], T]],
+    max_concurrency: int | None = None,
+) -> list[T] | list[ModelError]:
+    def _invoke_w_o_llm_error(invocable: Callable[[], T]) -> Result:
+        try:
+            return Result(value=invocable(), error=None)
+        except ModelError as e:
+            return Result(value=None, error=e)
+
+    with ThreadPoolExecutor(max_workers=max_concurrency) as executor:
+        results = list(executor.map(_invoke_w_o_llm_error, funcs))
+
+    errors: list[ModelError] = []
+    values = []
+    for res in results:
+        if res.error is not None:
+            errors.append(res.error)
+        else:
+            values.append(res.value)
+    if len(values) == 0:
+        assert len(errors) > 0
+        raise errors[0]
+    return values

chapter2/prompt-engineering/tau_bench/model_utils/api/logging.py

@@ -0,0 +1,74 @@
+import functools
+import inspect
+import json
+from multiprocessing import Lock
+from typing import Any
+
+from pydantic import BaseModel
+
+from tau_bench.model_utils.api.sample import SamplingStrategy
+from tau_bench.model_utils.model.utils import optionalize_type
+
+log_files = {}
+
+
+def prep_for_json_serialization(obj: Any, from_parse_method: bool = False):
+    # TODO: refine type annotations
+    if isinstance(obj, (str, int, float, bool, type(None))):
+        return obj
+    elif isinstance(obj, dict):
+        return {k: prep_for_json_serialization(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [prep_for_json_serialization(v) for v in obj]
+    elif isinstance(obj, tuple):
+        return tuple(prep_for_json_serialization(v) for v in obj)
+    elif isinstance(obj, set):
+        return {prep_for_json_serialization(v) for v in obj}
+    elif isinstance(obj, frozenset):
+        return frozenset(prep_for_json_serialization(v) for v in obj)
+    elif isinstance(obj, BaseModel):
+        return obj.model_dump(mode="json")
+    elif isinstance(obj, type) and issubclass(obj, BaseModel):
+        if from_parse_method:
+            optionalized_type = optionalize_type(obj)
+            return optionalized_type.model_json_schema()
+        else:
+            return obj.model_json_schema()
+    elif isinstance(obj, SamplingStrategy):
+        return obj.__class__.__name__
+    else:
+        raise TypeError(f"Object of type {type(obj)} is not JSON serializable")
+
+
+def log_call(func):
+    @functools.wraps(func)
+    def wrapper(self, *args, **kwargs):
+        response = func(self, *args, **kwargs)
+        log_file = getattr(self, "_log_file", None)
+        if log_file is not None:
+            if log_file not in log_files:
+                log_files[log_file] = Lock()
+            sig = inspect.signature(func)
+            bound_args = sig.bind(self, *args, **kwargs)
+            bound_args.apply_defaults()
+            all_args = bound_args.arguments
+            all_args.pop("self", None)
+
+            cls_name = self.__class__.__name__
+            log_entry = {
+                "cls_name": cls_name,
+                "method_name": func.__name__,
+                "kwargs": {
+                    k: prep_for_json_serialization(
+                        v, from_parse_method=func.__name__ in ["parse", "async_parse"]
+                    )
+                    for k, v in all_args.items()
+                },
+                "response": prep_for_json_serialization(response),
+            }
+            with log_files[log_file]:
+                with open(log_file, "a") as f:
+                    f.write(f"{json.dumps(log_entry)}\n")
+        return response
+
+    return wrapper

chapter2/prompt-engineering/tau_bench/model_utils/api/router.py

@@ -0,0 +1,92 @@
+import abc
+
+from pydantic import BaseModel
+
+from tau_bench.model_utils.api.datapoint import Datapoint, ScoreDatapoint
+from tau_bench.model_utils.model.model import Model
+
+
+class RequestRouter(abc.ABC):
+    @abc.abstractmethod
+    def route(self, dp: Datapoint, available_models: list[Model]) -> Model:
+        raise NotImplementedError
+
+
+class FirstModelRequestRouter(RequestRouter):
+    def route(self, dp: Datapoint, available_models: list[Model]) -> Model:
+        supporting_models = [model for model in available_models if model.supports_dp(dp)]
+        if len(supporting_models) == 0:
+            raise ValueError(f"No supporting models found from {available_models}")
+        return supporting_models[0]
+
+
+class CapabilityScoreModel(abc.ABC):
+    @abc.abstractmethod
+    def score_dp(self, dp: Datapoint) -> float:
+        raise NotImplementedError
+
+
+class PromptedLLMCapabilityScoreModel:
+    def __init__(self, model: Model | None = None) -> None:
+        if model is None:
+            from tau_bench.model_utils.model.claude import ClaudeModel
+
+            # claude is used as the default model as it is better at meta-level tasks
+            model = ClaudeModel()
+        self.model = model
+
+    def score_dp(self, dp: Datapoint, examples: list[ScoreDatapoint] | None = None) -> float:
+        return (
+            self.model.score(
+                instruction="Score the task in the datapoint on a scale of 1 (least complex) to 10 (most complex).",
+                text=f"----- start task -----\n{dp.model_dump_json()}\n----- end task -----",
+                min=1,
+                max=10,
+                examples=examples,
+            )
+            / 10.0
+        )
+
+
+class MinimumCapabilityRequestRouter(RequestRouter):
+    def __init__(self, capability_score_model: CapabilityScoreModel) -> None:
+        self.capability_score_model = capability_score_model
+
+    def route(self, dp: Datapoint, available_models: list[Model]) -> Model:
+        supporting_models = [model for model in available_models if model.supports_dp(dp)]
+        if len(supporting_models) == 0:
+            raise ValueError(f"No supporting models found from {available_models}")
+        required_capability = self.capability_score_model.score_dp(dp)
+        minimum_model: Model | None = None
+        minimum_model_capability: float | None = None
+        for model in supporting_models:
+            capability = model.get_capability()
+            if capability >= required_capability and (
+                minimum_model_capability is None or capability < minimum_model_capability
+            ):
+                minimum_model = model
+                minimum_model_capability = capability
+        if minimum_model is None:
+            raise ValueError(f"No model found with capability >= {required_capability}")
+        return minimum_model
+
+
+def request_router_factory(
+    router_id: str, capability_score_model: CapabilityScoreModel | None = None
+) -> RequestRouter:
+    if router_id == "first-model":
+        return FirstModelRequestRouter()
+    elif router_id == "minimum-capability":
+        if capability_score_model is None:
+            raise ValueError("CapabilityScoreModel is required for minimum-capability router")
+        return MinimumCapabilityRequestRouter(capability_score_model=capability_score_model)
+    raise ValueError(f"Unknown router_id: {router_id}")
+
+
+def default_request_router() -> RequestRouter:
+    return FirstModelRequestRouter()
+
+
+class RequestRouteDatapoint(BaseModel):
+    dp: Datapoint
+    capability_score: float

chapter2/prompt-engineering/tau_bench/model_utils/api/sample.py

@@ -0,0 +1,233 @@
+import abc
+import functools
+from multiprocessing import Lock
+from typing import Any, Callable, TypeVar
+
+from pydantic import BaseModel
+
+from tau_bench.model_utils.api.exception import APIError, execute_and_filter_model_errors
+from tau_bench.model_utils.model.exception import ModelError
+from tau_bench.model_utils import func_tools
+
+T = TypeVar("T")
+
+
+class SamplingStrategy(abc.ABC):
+    @abc.abstractmethod
+    def execute(self, invocable_or_invokables: Callable[..., T] | list[Callable[..., T]]) -> T:
+        raise NotImplementedError
+
+
+def catch_model_errors(func: Callable[..., T]) -> Callable[..., T]:
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs) -> T:
+        try:
+            return func(*args, **kwargs)
+        except ModelError as e:
+            raise APIError(
+                short_message=str(e),
+                report={
+                    "prompt": e.prompt,
+                    "response": e.response,
+                    "error_message": str(e),
+                },
+            )
+
+    return wrapper
+
+
+class SingleSamplingStrategy(SamplingStrategy):
+    @catch_model_errors
+    def execute(self, invocable_or_invokables: Callable[..., T]) -> T:
+        assert isinstance(invocable_or_invokables, Callable)
+        return invocable_or_invokables()
+
+
+class RedundantSamplingStrategy(SamplingStrategy):
+    def __init__(self, n: int = 2) -> None:
+        assert n > 0
+        self.n = n
+
+    @catch_model_errors
+    def execute(self, invocable_or_invokables: Callable[..., T] | list[Callable[..., T]]) -> T:
+        results = execute_and_filter_model_errors(
+            [lambda: invocable_or_invokables() for _ in range(self.n)]
+            if isinstance(invocable_or_invokables, Callable)
+            else invocable_or_invokables
+        )
+        assert len(results) > 0
+        return results[0]
+
+
+class RetrySamplingStrategy(SamplingStrategy):
+    def __init__(self, max_retries: int = 5) -> None:
+        assert max_retries > 0
+        self.max_retries = max_retries
+
+    @catch_model_errors
+    def execute(self, invocable_or_invokables: Callable[..., T]) -> T:
+        assert isinstance(invocable_or_invokables, Callable)
+        first_error = None
+        for _ in range(self.max_retries):
+            try:
+                return invocable_or_invokables()
+            except ModelError as e:
+                if first_error is None:
+                    first_error = e
+        assert first_error is not None
+        raise first_error
+
+
+class MajoritySamplingStrategy(SamplingStrategy):
+    def __init__(
+        self,
+        n: int = 5,
+        max_concurrency: int | None = None,
+        panic_on_first_model_error: bool = False,
+    ) -> None:
+        self.n = n
+        self.max_concurrency = max_concurrency if max_concurrency is not None else n
+        self.panic_on_first_model_error = panic_on_first_model_error
+
+    @catch_model_errors
+    def execute(self, invocable_or_invokables: Callable[..., T] | list[Callable[..., T]]) -> T:
+        if self.panic_on_first_model_error:
+            if isinstance(invocable_or_invokables, Callable):
+                results = list(
+                    func_tools.map(
+                        lambda _: invocable_or_invokables(),
+                        range(self.n),
+                        max_concurrency=self.max_concurrency,
+                    )
+                )
+            else:
+                results = list(
+                    func_tools.map(
+                        lambda invocable: invocable(),
+                        invocable_or_invokables,
+                        max_concurrency=self.max_concurrency,
+                    )
+                )
+        else:
+            results = execute_and_filter_model_errors(
+                (
+                    [lambda: invocable_or_invokables() for _ in range(self.n)]
+                    if isinstance(invocable_or_invokables, Callable)
+                    else invocable_or_invokables
+                ),
+                max_concurrency=self.max_concurrency,
+            )
+        if not self.panic_on_first_model_error and len(results) == 0:
+            raise SamplingError(
+                "No results from majority sampling (all calls resulted in LLM errors)"
+            )
+        return get_majority(results)
+
+
+def get_majority(results: list[T]) -> T:
+    grouped: dict[str, Any] = {}
+    for result in results:
+        if isinstance(result, BaseModel):
+            key = result.model_dump_json()
+        else:
+            key = str(result)
+        if key not in grouped:
+            # for now, just store duplicate results for the count
+            grouped[key] = [result]
+        else:
+            grouped[key].append(result)
+    majority = max(grouped, key=lambda key: len(grouped[key]))
+    return grouped[majority][0]
+
+
+class EnsembleSamplingStrategy(SamplingStrategy):
+    def __init__(
+        self, max_concurrency: int | None = None, panic_on_first_model_error: bool = False
+    ) -> None:
+        self.max_concurrency = max_concurrency
+        self.panic_on_first_model_error = panic_on_first_model_error
+
+    @catch_model_errors
+    def execute(self, invocable_or_invokables: Callable[..., T] | list[Callable[..., T]]) -> T:
+        if not isinstance(invocable_or_invokables, list) or len(invocable_or_invokables) < 2:
+            raise ValueError("Ensemble sampling requires at least 2 invocables")
+        if self.panic_on_first_model_error:
+            results = list(
+                func_tools.map(
+                    lambda invocable: invocable(),
+                    invocable_or_invokables,
+                    max_concurrency=self.max_concurrency,
+                )
+            )
+        else:
+            results = execute_and_filter_model_errors(
+                invocable_or_invokables, max_concurrency=self.max_concurrency
+            )
+        if not self.panic_on_first_model_error and len(results) == 0:
+            raise SamplingError(
+                "No results from ensemble sampling (all calls resulted in LLM errors)"
+            )
+        return get_majority(results)
+
+
+class UnanimousSamplingStrategy(SamplingStrategy):
+    def __init__(
+        self,
+        n: int = 5,
+        max_concurrency: int | None = None,
+        panic_on_first_model_error: bool = False,
+    ) -> None:
+        self.n = n
+        self.max_concurrency = max_concurrency if max_concurrency is not None else n
+        self.panic_on_first_model_error = panic_on_first_model_error
+
+    @catch_model_errors
+    def execute(self, invocable_or_invokables: Callable[..., T] | list[Callable[..., T]]) -> T:
+        if self.panic_on_first_model_error:
+            if isinstance(invocable_or_invokables, Callable):
+                results = list(
+                    func_tools.map(
+                        lambda _: invocable_or_invokables(),
+                        range(self.n),
+                        max_concurrency=self.max_concurrency,
+                    )
+                )
+            else:
+                results = list(
+                    func_tools.map(
+                        lambda invocable: invocable(),
+                        invocable_or_invokables,
+                        max_concurrency=self.max_concurrency,
+                    )
+                )
+        else:
+            results = execute_and_filter_model_errors(
+                (
+                    [lambda: invocable_or_invokables() for _ in range(self.n)]
+                    if isinstance(invocable_or_invokables, Callable)
+                    else invocable_or_invokables
+                ),
+                max_concurrency=self.max_concurrency,
+            )
+        if len(set(results)) > 1:
+            raise SamplingError("Results are not unanimous")
+        return results[0]
+
+
+class SamplingError(Exception):
+    pass
+
+
+DEFAULT_SAMPLING_STRATEGY = SingleSamplingStrategy()
+_DEFAULT_SAMPLING_STRATEGY_LOCK = Lock()
+
+
+def set_default_sampling_strategy(strategy: SamplingStrategy) -> None:
+    with _DEFAULT_SAMPLING_STRATEGY_LOCK:
+        global DEFAULT_SAMPLING_STRATEGY
+        DEFAULT_SAMPLING_STRATEGY = strategy
+
+
+def get_default_sampling_strategy() -> SamplingStrategy:
+    with _DEFAULT_SAMPLING_STRATEGY_LOCK:
+        return DEFAULT_SAMPLING_STRATEGY

chapter2/prompt-engineering/tau_bench/model_utils/api/tokens.py

@@ -0,0 +1,79 @@
+import json
+
+from pydantic import BaseModel
+
+from tau_bench.model_utils.api.datapoint import (
+    BinaryClassifyDatapoint,
+    ClassifyDatapoint,
+    Datapoint,
+    GenerateDatapoint,
+    ParseDatapoint,
+    ParseForceDatapoint,
+    ScoreDatapoint,
+)
+
+
+class TokenUsage(BaseModel):
+    input_tokens: int
+    output_tokens: int
+    by_primitive: dict[str, "TokenUsage"]
+
+
+def batch_token_analysis(dps: list[Datapoint], encoding_for_model: str = "gpt-4o") -> TokenUsage:
+    import tiktoken
+
+    enc = tiktoken.encoding_for_model(encoding_for_model)
+    # very rough estimates
+    inputs_by_primitive: dict[str, list[str]] = {}
+    outputs_by_primitive: dict[str, list[str]] = {}
+    for dp in dps:
+        input = json.dumps({k: v for k, v in dp.model_dump().items() if k != "response"})
+        inputs_by_primitive.setdefault(type(dp).__name__, []).append(input)
+        if isinstance(dp, ClassifyDatapoint):
+            output = f'{{"classification": {dp.response}}}'
+        elif isinstance(dp, BinaryClassifyDatapoint):
+            output = f'{{"classification": {0 if dp.response else 1}}}'
+        elif isinstance(dp, ParseForceDatapoint):
+            output = (
+                json.dumps(dp.response)
+                if isinstance(dp.response, dict)
+                else dp.response.model_dump_json()
+            )
+        elif isinstance(dp, GenerateDatapoint):
+            output = json.dumps(dp.response)
+        elif isinstance(dp, ParseDatapoint):
+            output = (
+                json.dumps(dp.response)
+                if isinstance(dp.response, dict)
+                else dp.response.model_dump_json()
+            )
+        elif isinstance(dp, ScoreDatapoint):
+            output = f"{{'score': {dp.response}}}"
+        else:
+            raise ValueError(f"Unknown datapoint type: {type(dp)}")
+        outputs_by_primitive.setdefault(type(dp).__name__, []).append(output)
+    input_tokens_by_primitive = {}
+    output_tokens_by_primitive = {}
+    for primitive, inputs in inputs_by_primitive.items():
+        input_tokens = sum([len(item) for item in enc.encode_batch(inputs)])
+        input_tokens_by_primitive[primitive] = input_tokens
+    for primitive, outputs in outputs_by_primitive.items():
+        output_tokens = sum([len(item) for item in enc.encode_batch(outputs)])
+        output_tokens_by_primitive[primitive] = output_tokens
+    return TokenUsage(
+        input_tokens=sum(input_tokens_by_primitive.values()),
+        output_tokens=sum(output_tokens_by_primitive.values()),
+        by_primitive={
+            primitive: TokenUsage(
+                input_tokens=input_tokens_by_primitive.get(primitive, 0),
+                output_tokens=output_tokens_by_primitive.get(primitive, 0),
+                by_primitive={},
+            )
+            for primitive in set(input_tokens_by_primitive.keys())
+            | set(output_tokens_by_primitive.keys())
+        },
+    )
+
+
+def token_analysis(dp: Datapoint, encoding_for_model: str = "gpt-4o") -> TokenUsage:
+    return batch_token_analysis([dp], encoding_for_model)

chapter2/prompt-engineering/tau_bench/model_utils/api/types.py

@@ -0,0 +1,3 @@
+from typing import Any
+
+PartialObj = dict[str, Any]

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

@@ -0,0 +1,11 @@
+import argparse
+
+from tau_bench.model_utils.model.model import Platform
+
+
+def api_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str)
+    parser.add_argument("--base-url", type=str)
+    parser.add_argument("--platform", type=str, required=True, choices=[e.value for e in Platform])
+    return parser

chapter2/prompt-engineering/tau_bench/model_utils/func_tools/__init__.py

@@ -0,0 +1,2 @@
+from tau_bench.model_utils.func_tools.filter import filter as filter
+from tau_bench.model_utils.func_tools.map import map as map

chapter2/prompt-engineering/tau_bench/model_utils/func_tools/filter.py

@@ -0,0 +1,17 @@
+from typing import Callable, Iterable, TypeVar
+
+from tau_bench.model_utils.func_tools.map import map
+
+T = TypeVar("T")
+
+builtin_filter = filter
+
+
+def filter(
+    func: Callable[[T], bool],
+    iterable: Iterable[T],
+    max_concurrency: int | None = None,
+) -> Iterable[T]:
+    assert max_concurrency is None or max_concurrency > 0
+    bits = map(func, iterable=iterable, max_concurrency=max_concurrency)
+    return [x for x, y in zip(iterable, bits) if y]

chapter2/prompt-engineering/tau_bench/model_utils/func_tools/map.py

@@ -0,0 +1,20 @@
+from concurrent.futures import ThreadPoolExecutor
+from typing import Callable, Iterable, TypeVar
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+def map(
+    func: Callable[[T], U],
+    iterable: Iterable[T],
+    max_concurrency: int | None = None,
+    use_tqdm: bool = False,
+) -> Iterable[U]:
+    assert max_concurrency is None or max_concurrency > 0
+    with ThreadPoolExecutor(max_workers=max_concurrency) as executor:
+        if use_tqdm:
+            from tqdm import tqdm
+
+            return list(tqdm(executor.map(func, iterable), total=len(iterable)))
+        return executor.map(func, iterable)

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)