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

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]