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

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
+)