cbrkit.synthesis.prompts

  1from collections.abc import Sequence
  2from dataclasses import dataclass, field
  3from typing import Any
  4
  5from ..dumpers import markdown
  6from ..helpers import get_value, sim_map2ranking, unpack_float, unpack_value
  7from ..typing import (
  8    Casebase,
  9    ConversionFunc,
 10    Float,
 11    JsonEntry,
 12    SimMap,
 13    StructuredValue,
 14    SynthesizerPromptFunc,
 15)
 16
 17__all__ = [
 18    "concat",
 19    "transpose",
 20    "transpose_value",
 21    "default",
 22    "pooling",
 23]
 24
 25
 26@dataclass(slots=True, frozen=True)
 27class concat[K, V, S: Float](SynthesizerPromptFunc[str, K, V, S]):
 28    """Concatenates multiple prompt functions into a single prompt."""
 29
 30    prompts: Sequence[SynthesizerPromptFunc[str, K, V, S] | str]
 31    separator: str = "\n\n"
 32
 33    def __call__(
 34        self,
 35        casebase: Casebase[K, V],
 36        query: V | None,
 37        similarities: SimMap[K, S] | None,
 38    ) -> str:
 39        return self.separator.join(
 40            prompt if isinstance(prompt, str) else prompt(casebase, query, similarities)
 41            for prompt in self.prompts
 42        )
 43
 44
 45@dataclass(slots=True, frozen=True)
 46class transpose[P, K, V1, V2, S: Float](SynthesizerPromptFunc[P, K, V1, S]):
 47    """Converts input values before passing them to a prompt function."""
 48
 49    prompt_func: SynthesizerPromptFunc[P, K, V2, S]
 50    conversion_func: ConversionFunc[V1, V2]
 51
 52    def __call__(
 53        self,
 54        casebase: Casebase[K, V1],
 55        query: V1 | None,
 56        similarities: SimMap[K, S] | None,
 57    ) -> P:
 58        return self.prompt_func(
 59            {key: self.conversion_func(value) for key, value in casebase.items()},
 60            self.conversion_func(query) if query is not None else None,
 61            similarities,
 62        )
 63
 64
 65def transpose_value[P, K, V, S: Float](
 66    func: SynthesizerPromptFunc[P, K, V, S],
 67) -> SynthesizerPromptFunc[P, K, StructuredValue[V], S]:
 68    """Wrap a prompt function to extract values from structured value inputs."""
 69    return transpose(func, get_value)
 70
 71
 72def encode[T](value: T, encoder: ConversionFunc[T, str]) -> str:
 73    """Encode a value to a string, using the encoder for non-primitive types."""
 74    if value is None:
 75        return ""
 76    elif isinstance(value, str):
 77        return value
 78    elif isinstance(value, int | float | bool):
 79        return str(value)
 80
 81    return encoder(value)
 82
 83
 84@dataclass(slots=True, frozen=True)
 85class default[V](SynthesizerPromptFunc[str, Any, V, Float]):
 86    """Produces an LLM input which provides context for the LLM to be able to perform instructions.
 87
 88    Args:
 89        instructions: Instructions for the LLM to execute on the input.
 90        encoder: Encoder function to convert the a case or query to a string.
 91        metadata: Optional metadata to include in the prompt.
 92
 93    Returns:
 94        A string to be used as an LLM input.
 95
 96    Examples:
 97        >>> prompt = default("Give me a summary of the found cars.")
 98        >>> result = prompt({"c1": "car A", "c2": "car B"}, "my query", {"c1": 0.9, "c2": 0.5})
 99        >>> "my query" in result
100        True
101    """
102
103    instructions: str | SynthesizerPromptFunc[str, Any, V, Float] | None = None
104    encoder: ConversionFunc[V | JsonEntry, str] = field(default_factory=markdown)
105    metadata: JsonEntry | None = None
106
107    def __call__(
108        self,
109        casebase: Casebase[Any, V],
110        query: V | None,
111        similarities: SimMap[Any, Float] | None,
112    ) -> str:
113        result = ""
114
115        if isinstance(self.instructions, str):
116            result += self.instructions
117        elif self.instructions is not None:
118            result += self.instructions(casebase, query, similarities)
119
120        if query is not None:
121            result += f"""
122## Query
123
124{encode(query, self.encoder)}
125"""
126
127        result += """
128## Documents Collection
129"""
130
131        ranking = (
132            sim_map2ranking(similarities)
133            if similarities is not None
134            else list(casebase.keys())
135        )
136
137        for rank, key in enumerate(ranking, start=1):
138            if similarities is not None:
139                result += f"""
140### {key} (Rank: {rank}, Similarity: {unpack_float(similarities[key]):.3f})
141"""
142            else:
143                result += f"""
144### {key}
145"""
146
147            result += f"""
148{encode(casebase[key], self.encoder)}
149"""
150
151        if self.metadata is not None:
152            result += f"""
153## Metadata
154
155{encode(self.metadata, self.encoder)}
156"""
157
158        return result
159
160
161@dataclass(slots=True, frozen=True)
162class pooling[V](ConversionFunc[Sequence[V], str]):
163    """
164    Produces an LLM input to aggregate partial results (i.e., the LLM output for single chunks) to a final, global result.
165
166    Args:
167        instructions: Instructions for the LLM to execute on the input.
168        encoder: Encoder function to convert the a case or query to a string.
169        metadata: Optional metadata to include in the prompt.
170
171    Examples:
172        >>> prompt = pooling("Please find the best match from the following partial results.")
173        >>> result = prompt(["result A", "result B", "result C"])
174        >>> "result A" in result
175        True
176    """
177
178    instructions: str | ConversionFunc[Sequence[V], str] | None = None
179    encoder: ConversionFunc[V | JsonEntry, str] = field(default_factory=markdown)
180    metadata: JsonEntry | None = None
181    unpack: bool = True
182
183    def __call__(
184        self,
185        values: Sequence[V],
186    ) -> str:
187        result = ""
188
189        if isinstance(self.instructions, str):
190            result += self.instructions
191        elif self.instructions is not None:
192            result += self.instructions(values)
193
194        result += """
195## Partial Results
196"""
197
198        for idx, value in enumerate(values, start=1):
199            if self.unpack:
200                value = unpack_value(value)
201
202            result += f"""
203### Result {idx}
204
205{encode(value, self.encoder)}
206"""
207
208        if self.metadata is not None:
209            result += f"""
210## Metadata
211
212{encode(self.metadata, self.encoder)}
213"""
214
215        return result