from langchain_core.prompts import StringPromptTemplate
from observability import span
from resolver import block

from blocks.base import BaseBlock


@block(name="Raw_Prompt", kind="llm")
class RawPromptChain(BaseBlock):
    """
    RawPromptChain render template with given text and pass the result to llm model.

    Example:

    ```
    node = RawPromptChain("Hello, {text}!")

    print(node("world"))
    ```

    Will output:

    "Hello, world!"
    """

    def __init__(self, prompt_template_type: StringPromptTemplate):
        self.prompt = prompt_template_type

    @span(name="RawPrompt")
    def __call__(self, text: str, **kwargs) -> str:
        return self.prompt.format(text=text, **kwargs)