cbrkit.synthesis.prompts

  1from collections.abc import Callable, 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)
 16from .providers.model import DocumentsPrompt
 17
 18__all__ = [
 19    "concat",
 20    "transpose",
 21    "transpose_value",
 22    "default",
 23    "documents_aware",
 24    "pooling",
 25]
 26
 27
 28@dataclass(slots=True, frozen=True)
 29class concat[K, V, S: Float](SynthesizerPromptFunc[str, K, V, S]):
 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    prompt_func: SynthesizerPromptFunc[P, K, V2, S]
 48    conversion_func: ConversionFunc[V1, V2]
 49
 50    def __call__(
 51        self,
 52        casebase: Casebase[K, V1],
 53        query: V1 | None,
 54        similarities: SimMap[K, S] | None,
 55    ) -> P:
 56        return self.prompt_func(
 57            {key: self.conversion_func(value) for key, value in casebase.items()},
 58            self.conversion_func(query) if query is not None else None,
 59            similarities,
 60        )
 61
 62
 63def transpose_value[P, K, V, S: Float](
 64    func: SynthesizerPromptFunc[P, K, V, S],
 65) -> SynthesizerPromptFunc[P, K, StructuredValue[V], S]:
 66    return transpose(func, get_value)
 67
 68
 69def encode[T](value: T, encoder: ConversionFunc[T, str]) -> str:
 70    if value is None:
 71        return ""
 72    elif isinstance(value, str):
 73        return value
 74    elif isinstance(value, int | float | bool):
 75        return str(value)
 76
 77    return encoder(value)
 78
 79
 80@dataclass(slots=True, frozen=True)
 81class default[V](SynthesizerPromptFunc[str, Any, V, Float]):
 82    """Produces an LLM input which provides context for the LLM to be able to perform instructions.
 83
 84    Args:
 85        instructions: Instructions for the LLM to execute on the input.
 86        encoder: Encoder function to convert the a case or query to a string.
 87        metadata: Optional metadata to include in the prompt.
 88
 89    Returns:
 90        A string to be used as an LLM input.
 91
 92    Examples:
 93        >>> prompt = default("Give me a summary of the found cars.")
 94        >>> prompt(casebase, query, similarities) # doctest: +SKIP
 95    """
 96
 97    instructions: str | SynthesizerPromptFunc[str, Any, V, Float] | None = None
 98    encoder: ConversionFunc[V | JsonEntry, str] = field(default_factory=markdown)
 99    metadata: JsonEntry | None = None
100
101    def __call__(
102        self,
103        casebase: Casebase[Any, V],
104        query: V | None,
105        similarities: SimMap[Any, Float] | None,
106    ) -> str:
107        result = ""
108
109        if isinstance(self.instructions, Callable):
110            result += self.instructions(casebase, query, similarities)
111        elif isinstance(self.instructions, str):
112            result += self.instructions
113
114        if query is not None:
115            result += f"""
116## Query
117
118{encode(query, self.encoder)}
119"""
120
121        result += """
122## Documents Collection
123"""
124
125        ranking = (
126            sim_map2ranking(similarities)
127            if similarities is not None
128            else list(casebase.keys())
129        )
130
131        for rank, key in enumerate(ranking, start=1):
132            if similarities is not None:
133                result += f"""
134### {key} (Rank: {rank}, Similarity: {unpack_float(similarities[key]):.3f})
135"""
136            else:
137                result += f"""
138### {key}
139"""
140
141            result += f"""
142{encode(casebase[key], self.encoder)}
143"""
144
145        if self.metadata is not None:
146            result += f"""
147## Metadata
148
149{encode(self.metadata, self.encoder)}
150"""
151
152        return result
153
154
155@dataclass(slots=True, frozen=True)
156class documents_aware[V](SynthesizerPromptFunc[DocumentsPrompt[str], Any, V, Any]):
157    """
158    Produces a structured LLM input (as of now: exclusive for cohere) which provides context for the LLM to be able to perform instructions.
159
160    Args:
161        instructions: Instructions for the LLM to execute on the input.
162        encoder: Encoder function to convert the a case or query to a string.
163        metadata: Optional metadata to include in the prompt.
164
165    Examples:
166        >>> prompt = documents_aware("Give me a summary of the found cars.")
167        >>> prompt(casebase, query, similarities) # doctest: +SKIP
168    """
169
170    instructions: str | SynthesizerPromptFunc[str, Any, V, Float] | None = None
171    encoder: ConversionFunc[V | JsonEntry, str] = field(default_factory=markdown)
172    metadata: JsonEntry | None = None
173
174    def __call__(
175        self,
176        casebase: Casebase[Any, V],
177        query: V | None,
178        similarities: SimMap[Any, Float] | None,
179    ) -> DocumentsPrompt:
180        result = ""
181
182        if isinstance(self.instructions, Callable):
183            result += self.instructions(casebase, query, similarities)
184        elif isinstance(self.instructions, str):
185            result += self.instructions
186
187        if query is not None:
188            result += f"""
189## Query
190
191{encode(query, self.encoder)}
192"""
193
194        if self.metadata is not None:
195            result += f"""
196## Metadata
197
198{encode(self.metadata, self.encoder)}
199"""
200
201        ranking = (
202            sim_map2ranking(similarities)
203            if similarities is not None
204            else list(casebase.keys())
205        )
206
207        return DocumentsPrompt(
208            result,
209            {
210                key: {
211                    "text": encode(casebase[key], self.encoder),
212                    "similarity": f"{unpack_float(similarities[key]):.3f}",
213                    "rank": str(rank),
214                }
215                if similarities is not None
216                else {
217                    "text": encode(casebase[key], self.encoder),
218                }
219                for rank, key in enumerate(ranking)
220            },
221        )
222
223
224@dataclass(slots=True, frozen=True)
225class pooling[V](ConversionFunc[Sequence[V], str]):
226    """
227    Produces an LLM input to aggregate partial results (i.e., the LLM output for single chunks) to a final, global result.
228
229    Args:
230        instructions: Instructions for the LLM to execute on the input.
231        encoder: Encoder function to convert the a case or query to a string.
232        metadata: Optional metadata to include in the prompt.
233
234    Examples:
235        >>> prompt = pooling("Please find the best match from the following partial results.")
236        >>> prompt([partial1, partial2, partial3]) # doctest: +SKIP
237    """
238
239    instructions: str | ConversionFunc[Sequence[V], str] | None = None
240    encoder: ConversionFunc[V | JsonEntry, str] = field(default_factory=markdown)
241    metadata: JsonEntry | None = None
242    unpack: bool = True
243
244    def __call__(
245        self,
246        values: Sequence[V],
247    ) -> str:
248        result = ""
249
250        if isinstance(self.instructions, Callable):
251            result += self.instructions(values)
252        elif isinstance(self.instructions, str):
253            result += self.instructions
254
255        result += """
256## Partial Results
257"""
258
259        for idx, value in enumerate(values, start=1):
260            if self.unpack:
261                value = unpack_value(value)
262
263            result += f"""
264### Result {idx}
265
266{encode(value, self.encoder)}
267"""
268
269        if self.metadata is not None:
270            result += f"""
271## Metadata
272
273{encode(self.metadata, self.encoder)}
274"""
275
276        return result