cbrkit.synthesis.providers

34@dataclass(slots=True, kw_only=True)
35class AsyncProvider[P, R](BatchConversionFunc[P, R], ABC):
36    """Base class for async batch-processing providers."""
37
38    def __call__(self, batches: Sequence[P]) -> Sequence[R]:
39        return run_coroutine(self.__call_batches__(batches))
40
41    async def __call_batches__(self, batches: Sequence[P]) -> Sequence[R]:
42        logger.info(f"Processing {len(batches)} batches")
43
44        return await asyncio.gather(
45            *(
46                self.__call_batch_wrapper__(batch, idx)
47                for idx, batch in enumerate(batches)
48            )
49        )
50
51    async def __call_batch_wrapper__(self, prompt: P, idx: int) -> R:
52        result = await self.__call_batch__(prompt)
53        logger.debug(f"Result of batch {idx + 1}: {result}")
54        return result
55
56    @abstractmethod
57    async def __call_batch__(self, prompt: P) -> R: ...

60@dataclass(slots=True, kw_only=True)
61class BaseProvider[P, R](AsyncProvider[P, Response[R]], ABC):
62    """Base provider with model configuration, retry logic, and error handling."""
63
64    model: str
65    response_type: type[R]
66    default_response: R | None = None
67    system_message: str | None = None
68    delay: float = 0
69    retries: int = 0
70    extra_kwargs: Mapping[str, Any] = field(default_factory=dict)
71
72    @override
73    async def __call_batch_wrapper__(
74        self, prompt: P, idx: int, retry: int = 0
75    ) -> Response[R]:
76        if self.delay > 0 and retry == 0:
77            await asyncio.sleep(idx * self.delay)
78
79        try:
80            return await super(BaseProvider, self).__call_batch_wrapper__(prompt, idx)
81
82        except Exception as e:
83            if retry < self.retries:
84                logger.info(f"Retrying batch {idx + 1}...")
85                return await self.__call_batch_wrapper__(prompt, idx, retry + 1)
86
87            if self.default_response is not None:
88                logger.error(f"Error processing batch {idx + 1}: {e}")
89                return Response(self.default_response, Usage(0, 0))
90
91            raise

27@dataclass(slots=True, frozen=True)
28class Response[T](StructuredValue[T]):
29    """Provider response wrapping a value with usage statistics."""
30
31    usage: Usage = field(default_factory=Usage)

14@dataclass(slots=True, frozen=True)
15class Usage:
16    """Token usage statistics for a provider response."""
17
18    prompt_tokens: int = 0
19    completion_tokens: int = 0
20
21    @property
22    def total_tokens(self) -> int:
23        """Return the sum of prompt and completion tokens."""
24        return self.prompt_tokens + self.completion_tokens

 28    @dataclass(slots=True)
 29    class anthropic[R: str | BaseModel](BaseProvider[AnthropicPrompt, R]):
 30        """Provider that calls Anthropic's messages API."""
 31
 32        model: ModelParam
 33        max_tokens: int
 34        messages: Sequence[MessageParam] = field(default_factory=tuple)
 35        client: AsyncAnthropic = field(default_factory=AsyncAnthropic, repr=False)
 36        metadata: MetadataParam | None = None
 37        stop_sequences: list[str] | None = None
 38        tool_choice: ToolChoiceParam | None = None
 39        tools: Iterable[ToolUnionParam] | None = None
 40        extra_headers: Any | None = None
 41        extra_query: Any | None = None
 42        extra_body: Any | None = None
 43        timeout: float | Timeout | None | NotGiven = NOT_GIVEN
 44
 45        @override
 46        async def __call_batch__(self, prompt: AnthropicPrompt) -> Response[R]:
 47            messages: list[MessageParam] = []
 48
 49            messages.extend(self.messages)
 50
 51            if isinstance(prompt, str):
 52                messages.append({"role": "user", "content": prompt})
 53            else:
 54                messages.extend(prompt)
 55
 56            output_format = (
 57                cast(Any, self.response_type)
 58                if issubclass(self.response_type, BaseModel)
 59                else omit
 60            )
 61            res = await self.client.messages.parse(
 62                model=self.model,
 63                messages=messages,
 64                max_tokens=self.max_tokens,
 65                output_format=output_format,
 66                metadata=if_given(self.metadata),
 67                stop_sequences=if_given(self.stop_sequences),
 68                system=if_given(self.system_message),
 69                tool_choice=if_given(self.tool_choice),
 70                tools=if_given(self.tools),
 71                extra_headers=self.extra_headers,
 72                extra_query=self.extra_query,
 73                extra_body=self.extra_body,
 74                timeout=self.timeout,
 75                **self.extra_kwargs,
 76            )
 77
 78            usage = Usage(
 79                res.usage.input_tokens,
 80                res.usage.output_tokens,
 81            )
 82
 83            if (
 84                isinstance(self.response_type, type)
 85                and issubclass(self.response_type, BaseModel)
 86                and (parsed := res.parsed_output) is not None
 87            ):
 88                return Response(parsed, usage)
 89
 90            if (
 91                isinstance(self.response_type, type)
 92                and issubclass(self.response_type, str)
 93                and len(res.content) > 0
 94            ):
 95                aggregated_content = "".join(
 96                    getattr(block, "text", "") for block in res.content
 97                )
 98                return Response(cast(R, aggregated_content), usage)
 99
100            raise ValueError("Invalid response", res)

 33    @dataclass(slots=True)
 34    class cohere[R: str | BaseModel](BaseProvider[CoherePrompt, R]):
 35        """Provider that calls Cohere's chat API."""
 36
 37        messages: Sequence[ChatMessageV2] = field(default_factory=tuple)
 38        documents: Sequence[V2ChatRequestDocumentsItem] = field(default_factory=tuple)
 39        client: AsyncClient = field(default_factory=AsyncClient, repr=False)
 40        request_options: RequestOptions | None = None
 41        citation_options: CitationOptions | None = None
 42        safety_mode: V2ChatRequestSafetyMode | None = None
 43        max_tokens: int | None = None
 44        stop_sequences: Sequence[str] | None = None
 45        temperature: float | None = None
 46        seed: int | None = None
 47        frequency_penalty: float | None = None
 48        presence_penalty: float | None = None
 49        k: int | None = None
 50        p: float | None = None
 51        logprobs: bool | None = None
 52
 53        @override
 54        async def __call_batch__(self, prompt: CoherePrompt) -> Response[R]:
 55            documents: list[V2ChatRequestDocumentsItem] = list(self.documents)
 56
 57            if isinstance(prompt, CohereDocumentsPrompt):
 58                documents.extend(prompt.documents)
 59
 60            if issubclass(self.response_type, BaseModel) and documents:
 61                raise ValueError(
 62                    "Structured output format is not supported when using documents"
 63                )
 64
 65            messages: list[ChatMessageV2] = []
 66
 67            if self.system_message is not None:
 68                messages.append(SystemChatMessageV2(content=self.system_message))
 69
 70            if isinstance(prompt, str):
 71                messages.append(UserChatMessageV2(content=prompt))
 72            elif isinstance(prompt, CohereDocumentsPrompt):
 73                messages.extend(prompt.messages)
 74            else:
 75                messages.extend(prompt)
 76
 77            res = await self.client.v2.chat(
 78                model=self.model,
 79                messages=messages,
 80                request_options=self.request_options,
 81                documents=documents if documents else None,
 82                response_format=JsonObjectResponseFormatV2(
 83                    json_schema=self.response_type.model_json_schema()
 84                )
 85                if issubclass(self.response_type, BaseModel)
 86                else None,
 87                citation_options=self.citation_options,
 88                safety_mode=self.safety_mode,
 89                max_tokens=self.max_tokens,
 90                stop_sequences=self.stop_sequences,
 91                temperature=self.temperature,
 92                seed=self.seed,
 93                frequency_penalty=self.frequency_penalty,
 94                presence_penalty=self.presence_penalty,
 95                k=self.k,
 96                p=self.p,
 97                logprobs=self.logprobs,
 98                **self.extra_kwargs,
 99            )
100
101            content = res.message.content
102
103            if content is None:
104                raise ValueError("The completion is empty")
105
106            if issubclass(self.response_type, BaseModel):
107                if len(content) != 1 or content[0].type != "text":
108                    raise ValueError(
109                        "The completion is empty, has multiple outputs, or is not text"
110                    )
111
112                return Response(
113                    cast(R, self.response_type.model_validate_json(content[0].text))
114                )
115
116            aggregated_content = "".join(
117                block.text for block in content if block.type == "text"
118            )
119
120            return Response(cast(R, aggregated_content))

18@dataclass(slots=True, frozen=True)
19class conversation[P, R](ConversionFunc[Sequence[P], R]):
20    """Iteratively generates responses until the conversion function returns None."""
21
22    generation_func: AnyConversionFunc[Sequence[P], R]
23    conversion_func: ConversionFunc[R, Sequence[P] | None]
24
25    def __call__(self, batch: Sequence[P]) -> R:
26        func = unbatchify_conversion(self.generation_func)
27        result = func(batch)
28
29        while next_batch := self.conversion_func(result):
30            result = func(next_batch)
31
32        return result

20    @dataclass(slots=True)
21    class google[R: BaseModel | str](BaseProvider[GooglePrompt, R]):
22        """Provider that calls Google's Generative AI API."""
23
24        client: Client = field(default_factory=Client, repr=False)
25        config: GenerateContentConfig = field(init=False)
26        base_config: InitVar[GenerateContentConfig | None] = None
27
28        def __post_init__(self, base_config: GenerateContentConfig | None) -> None:
29            self.config = base_config or GenerateContentConfig()
30
31            if issubclass(self.response_type, BaseModel):
32                self.config.response_schema = self.response_type
33
34            if self.system_message is not None:
35                self.config.system_instruction = self.system_message
36
37        @override
38        async def __call_batch__(self, prompt: GooglePrompt) -> Response[R]:
39            res = await self.client.aio.models.generate_content(
40                model=self.model,
41                contents=prompt,
42                config=self.config,
43                **self.extra_kwargs,
44            )
45
46            if (
47                issubclass(self.response_type, BaseModel)
48                and (parsed := res.parsed)
49                and isinstance(parsed, self.response_type)
50            ):
51                return Response(cast(R, parsed))
52
53            elif issubclass(self.response_type, str) and (text := res.text):
54                return Response(cast(R, text))
55
56            raise ValueError("Invalid response", res)

17    @dataclass(slots=True)
18    class ollama[R: str | BaseModel](BaseProvider[OllamaPrompt, R]):
19        """Provider that calls Ollama's chat API."""
20
21        client: AsyncClient = field(default_factory=AsyncClient, repr=False)
22        messages: Sequence[Message] = field(default_factory=tuple)
23        options: Options | None = None
24        keep_alive: float | str | None = None
25
26        @override
27        async def __call_batch__(self, prompt: OllamaPrompt) -> Response[R]:
28            messages: list[Message] = []
29
30            if self.system_message is not None:
31                messages.append(Message(role="system", content=self.system_message))
32
33            messages.extend(self.messages)
34
35            if isinstance(prompt, str):
36                messages.append(Message(role="user", content=prompt))
37            else:
38                messages.extend(prompt)
39
40            res = await self.client.chat(
41                model=self.model,
42                messages=messages,
43                options=self.options,
44                keep_alive=self.keep_alive,
45                format=self.response_type.model_json_schema()
46                if issubclass(self.response_type, BaseModel)
47                else None,
48                **self.extra_kwargs,
49            )
50
51            content = res["message"]["content"]
52
53            if self.response_type is str:
54                return Response(content)
55
56            return Response(json.loads(content))

33    @dataclass(slots=True)
34    class openai_agents[T, R](AsyncProvider[OpenaiAgentsPrompt, TypedRunResult[R]]):
35        """Provider that runs OpenAI Agents SDK agents."""
36
37        agents: MaybeSequence[Agent[T]]
38        context: T | None = None
39        max_turns: int = DEFAULT_MAX_TURNS
40        hooks: RunHooks[T] | None = None
41        run_config: RunConfig | None = None
42
43        @override
44        async def __call_batch__(self, prompt: OpenaiAgentsPrompt) -> TypedRunResult[R]:
45            agents = produce_sequence(self.agents)
46
47            if not agents:
48                raise ValueError("No agents given.")
49
50            head_agent, *tail_agents = agents
51
52            session: Any = SQLiteSession(uuid1().hex) if len(agents) > 1 else None
53
54            run = partial(
55                Runner.run,
56                context=self.context,
57                max_turns=self.max_turns,
58                hooks=self.hooks,
59                run_config=self.run_config,
60                session=session,
61            )
62
63            response: RunResult = await run(head_agent, prompt)
64
65            for agent in tail_agents:
66                response = await run(agent, [])
67
68            return cast(TypedRunResult[R], response)

 37    @dataclass(slots=True)
 38    class openai_completions[R: BaseModel | str](BaseProvider[OpenAiPrompt, R]):
 39        """Provider that calls OpenAI's chat completions API."""
 40
 41        model: str | ChatModel
 42        messages: Sequence[ChatCompletionMessageParam] = field(default_factory=tuple)
 43        tool_choice: type[BaseModel] | str | None = None
 44        client: AsyncOpenAI = field(default_factory=AsyncOpenAI, repr=False)
 45        frequency_penalty: float | None = None
 46        logit_bias: dict[str, int] | None = None
 47        logprobs: bool | None = None
 48        max_completion_tokens: int | None = None
 49        metadata: dict[str, str] | None = None
 50        n: int | None = None
 51        presence_penalty: float | None = None
 52        seed: int | None = None
 53        stop: str | list[str] | None = None
 54        store: bool | None = None
 55        reasoning_effort: Literal["low", "medium", "high"] | None = None
 56        temperature: float | None = None
 57        top_logprobs: int | None = None
 58        top_p: float | None = None
 59        extra_headers: Any | None = None
 60        extra_query: Any | None = None
 61        extra_body: Any | None = None
 62        timeout: float | Timeout | None | NotGiven = NOT_GIVEN
 63
 64        @override
 65        async def __call_batch__(self, prompt: OpenAiPrompt) -> Response[R]:
 66            messages: list[ChatCompletionMessageParam] = []
 67
 68            if self.system_message is not None:
 69                messages.append({"role": "system", "content": self.system_message})
 70
 71            messages.extend(self.messages)
 72
 73            if isinstance(prompt, str):
 74                messages.append({"role": "user", "content": prompt})
 75            else:
 76                messages.extend(prompt)
 77
 78            tools: list[ChatCompletionFunctionToolParam] | None = None
 79            tool_choice: ChatCompletionNamedToolChoiceParam | None = None
 80            response_type_origin = get_origin(self.response_type)
 81
 82            if response_type_origin is UnionType or response_type_origin is Union:
 83                tools = [
 84                    pydantic_function_tool(tool)
 85                    for tool in get_args(self.response_type)
 86                    if issubclass(tool, BaseModel)
 87                ]
 88            elif (
 89                issubclass(self.response_type, BaseModel)
 90                and self.tool_choice is not None
 91            ):
 92                tools = [pydantic_function_tool(self.response_type)]
 93
 94            if self.tool_choice is not None:
 95                tool_choice = {
 96                    "type": "function",
 97                    "function": {
 98                        "name": self.tool_choice
 99                        if isinstance(self.tool_choice, str)
100                        else self.tool_choice.__name__,
101                    },
102                }
103
104            response_format = (
105                cast(Any, self.response_type)
106                if tools is None and issubclass(self.response_type, BaseModel)
107                else omit
108            )
109
110            try:
111                res = await self.client.chat.completions.parse(
112                    model=self.model,
113                    messages=messages,
114                    response_format=response_format,
115                    tools=if_given(tools),
116                    tool_choice=if_given(tool_choice),
117                    frequency_penalty=if_given(self.frequency_penalty),
118                    logit_bias=if_given(self.logit_bias),
119                    logprobs=if_given(self.logprobs),
120                    max_completion_tokens=if_given(self.max_completion_tokens),
121                    metadata=if_given(self.metadata),
122                    n=if_given(self.n),
123                    presence_penalty=if_given(self.presence_penalty),
124                    seed=if_given(self.seed),
125                    stop=if_given(self.stop),
126                    store=if_given(self.store),
127                    reasoning_effort=if_given(self.reasoning_effort),
128                    temperature=if_given(self.temperature),
129                    top_logprobs=if_given(self.top_logprobs),
130                    top_p=if_given(self.top_p),
131                    extra_headers=self.extra_headers,
132                    extra_query=self.extra_query,
133                    extra_body=self.extra_body,
134                    timeout=self.timeout,
135                    **self.extra_kwargs,
136                )
137            except ValidationError as e:
138                for error in e.errors():
139                    logger.error(f"Invalid response ({error['msg']}): {error['input']}")
140                raise
141
142            choice = res.choices[0]
143            message = choice.message
144
145            assert res.usage is not None
146            usage = Usage(res.usage.prompt_tokens, res.usage.completion_tokens)
147
148            if choice.finish_reason == "length":
149                raise ValueError("Length limit", res)
150
151            if choice.finish_reason == "content_filter":
152                raise ValueError("Content filter", res)
153
154            if message.refusal:
155                raise ValueError("Refusal", res)
156
157            if (
158                isinstance(self.response_type, type)
159                and issubclass(self.response_type, BaseModel)
160                and (parsed := message.parsed) is not None
161            ):
162                return Response(cast(R, parsed), usage)
163
164            if (
165                isinstance(self.response_type, type)
166                and issubclass(self.response_type, str)
167                and (content := message.content) is not None
168            ):
169                return Response(cast(R, content), usage)
170
171            if (
172                tools is not None
173                and (tool_calls := message.tool_calls) is not None
174                and (parsed := tool_calls[0].function.parsed_arguments) is not None
175            ):
176                return Response(cast(R, parsed), usage)
177
178            raise ValueError("Invalid response", res)

 40    @dataclass(slots=True)
 41    class openai_responses[R: BaseModel | str](BaseProvider[OpenAiResponsesPrompt, R]):
 42        """Provider that calls the OpenAI Responses API and parses structured outputs."""
 43
 44        input_items: Sequence[ResponseInputItemParam] = field(default_factory=tuple)
 45        tool_choice: type[BaseModel] | str | None = None
 46        client: AsyncOpenAI = field(default_factory=AsyncOpenAI, repr=False)
 47        include: Sequence[ResponseIncludable] | None = None
 48        max_output_tokens: int | None = None
 49        max_tool_calls: int | None = None
 50        metadata: dict[str, str] | None = None
 51        parallel_tool_calls: bool | None = None
 52        store: bool | None = None
 53        temperature: float | None = None
 54        top_logprobs: int | None = None
 55        top_p: float | None = None
 56        text: ResponseTextConfigParam | None = None
 57        extra_headers: Any | None = None
 58        extra_query: Any | None = None
 59        extra_body: Any | None = None
 60        timeout: float | Timeout | None | NotGiven = NOT_GIVEN
 61
 62        @override
 63        async def __call_batch__(self, prompt: OpenAiResponsesPrompt) -> Response[R]:
 64            inputs: list[ResponseInputItemParam] = []
 65
 66            if self.system_message is not None:
 67                inputs.append({"role": "system", "content": self.system_message})
 68
 69            inputs.extend(self.input_items)
 70
 71            if isinstance(prompt, str):
 72                inputs.append({"role": "user", "content": prompt})
 73            else:
 74                inputs.extend(prompt)
 75
 76            tools: list[ParseableToolParam] | None = None
 77            tool_choice: ToolChoice | None = None
 78            text_format: type[BaseModel] | Omit = omit
 79
 80            response_type_origin = get_origin(self.response_type)
 81
 82            if response_type_origin is UnionType or response_type_origin is Union:
 83                tools = [
 84                    cast(ParseableToolParam, pydantic_function_tool(tool))
 85                    for tool in get_args(self.response_type)
 86                    if issubclass(tool, BaseModel)
 87                ]
 88            elif issubclass(self.response_type, BaseModel):
 89                if self.tool_choice is not None:
 90                    tools = [
 91                        cast(
 92                            ParseableToolParam,
 93                            pydantic_function_tool(self.response_type),
 94                        )
 95                    ]
 96                else:
 97                    text_format = self.response_type
 98
 99            if self.tool_choice is not None:
100                tool_choice = ToolChoiceFunctionParam(
101                    name=self.tool_choice
102                    if isinstance(self.tool_choice, str)
103                    else self.response_type.__name__,
104                    type="function",
105                )
106
107            text_param: ResponseTextConfigParam | Omit
108
109            if self.text is None:
110                text_param = omit
111            elif text_format is not omit and "format" in self.text:
112                raise ValueError(
113                    "`text.format` cannot be set when using structured outputs."
114                )
115            else:
116                text_param = self.text
117
118            try:
119                res = await self.client.responses.parse(
120                    model=self.model,
121                    input=inputs,
122                    instructions=if_given(self.system_message),
123                    include=if_given(
124                        list(self.include) if self.include is not None else None
125                    ),
126                    tools=if_given(tools),
127                    tool_choice=if_given(tool_choice),
128                    max_output_tokens=if_given(self.max_output_tokens),
129                    max_tool_calls=if_given(self.max_tool_calls),
130                    metadata=if_given(self.metadata),
131                    parallel_tool_calls=if_given(self.parallel_tool_calls),
132                    store=if_given(self.store),
133                    temperature=if_given(self.temperature),
134                    top_logprobs=if_given(self.top_logprobs),
135                    top_p=if_given(self.top_p),
136                    text=text_param,
137                    text_format=cast(Any, text_format),
138                    extra_headers=self.extra_headers,
139                    extra_query=self.extra_query,
140                    extra_body=self.extra_body,
141                    timeout=self.timeout,
142                    **self.extra_kwargs,
143                )
144            except ValidationError as e:
145                for error in e.errors():
146                    logger.error(f"Invalid response ({error['msg']}): {error['input']}")
147                raise
148
149            if res.incomplete_details is not None:
150                raise ValueError(
151                    res.incomplete_details.reason or "Response incomplete", res
152                )
153
154            for output in res.output:
155                content_list = getattr(output, "content", None)
156                if content_list is not None:
157                    for content in content_list:
158                        if content.type == "refusal":
159                            raise ValueError("Refusal", res)
160
161            assert res.usage is not None
162            usage = Usage(res.usage.input_tokens, res.usage.output_tokens)
163
164            if tools is not None:
165                for output in res.output:
166                    if output.type == "function_call":
167                        parsed_arguments = getattr(output, "parsed_arguments", None)
168
169                        if parsed_arguments is not None:
170                            return Response(cast(R, parsed_arguments), usage)
171
172                raise ValueError("Invalid response", res)
173
174            if text_format is not omit and (parsed := res.output_parsed) is not None:
175                return Response(cast(R, parsed), usage)
176
177            if issubclass(self.response_type, str):
178                content = res.output_text
179
180                if content:
181                    return Response(cast(R, content), usage)
182
183            raise ValueError("Invalid response", res)