File size: 826 Bytes
719d0db
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
from abc import ABC, abstractmethod
from typing import Dict, Any 
from langchain_core.runnables import RunnableLambda
from langchain_core.runnables.base import Runnable

class TemplateJsonBase(ABC):
    parser: Runnable
    template: str
    prompt: Runnable

    @abstractmethod
    def _get_output_key(self) -> str:
        raise NotImplementedError

    def get_template(self) -> str:
        return self.template
    
    def extract_value(self, input: Dict[str, Any]) -> Any:
        return input[self._get_output_key()]

    def sandwiches(self, 
                   llm: Runnable,
                   extract_value: bool = False) -> Runnable:
        if extract_value:
            return self.prompt | llm | self.parser | RunnableLambda(self.extract_value)
        else:
            return self.prompt | llm | self.parser