o2o-infinith-backend/app/integrations/llm/prompt.py

20 lines
504 B
Python

from pydantic import BaseModel
class Prompt:
def __init__(
self,
template: str,
model: str,
input_class: type[BaseModel],
output_class: type[BaseModel],
):
self.template = template
self.model = model
self.input_class = input_class
self.output_class = output_class
def build(self, input_data: dict) -> str:
verified = self.input_class(**input_data)
return self.template.format(**verified.model_dump())