File size: 3,178 Bytes
129cd69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
from __future__ import annotations

from typing import Any, Callable, Dict, Optional

from langchain_core.output_parsers import BaseOutputParser


class GuardrailsOutputParser(BaseOutputParser):
    """Parse the output of an LLM call using Guardrails."""

    guard: Any
    """The Guardrails object."""
    api: Optional[Callable]
    """The API to use for the Guardrails object."""
    args: Any
    """The arguments to pass to the API."""
    kwargs: Any
    """The keyword arguments to pass to the API."""

    @property
    def _type(self) -> str:
        return "guardrails"

    @classmethod
    def from_rail(
        cls,
        rail_file: str,
        num_reasks: int = 1,
        api: Optional[Callable] = None,
        *args: Any,
        **kwargs: Any,
    ) -> GuardrailsOutputParser:
        """Create a GuardrailsOutputParser from a rail file.

        Args:
            rail_file: a rail file.
            num_reasks: number of times to re-ask the question.
            api: the API to use for the Guardrails object.
            *args: The arguments to pass to the API
            **kwargs: The keyword arguments to pass to the API.

        Returns:
            GuardrailsOutputParser
        """
        try:
            from guardrails import Guard
        except ImportError:
            raise ImportError(
                "guardrails-ai package not installed. "
                "Install it by running `pip install guardrails-ai`."
            )
        return cls(
            guard=Guard.from_rail(rail_file, num_reasks=num_reasks),
            api=api,
            args=args,
            kwargs=kwargs,
        )

    @classmethod
    def from_rail_string(
        cls,
        rail_str: str,
        num_reasks: int = 1,
        api: Optional[Callable] = None,
        *args: Any,
        **kwargs: Any,
    ) -> GuardrailsOutputParser:
        try:
            from guardrails import Guard
        except ImportError:
            raise ImportError(
                "guardrails-ai package not installed. "
                "Install it by running `pip install guardrails-ai`."
            )
        return cls(
            guard=Guard.from_rail_string(rail_str, num_reasks=num_reasks),
            api=api,
            args=args,
            kwargs=kwargs,
        )

    @classmethod
    def from_pydantic(
        cls,
        output_class: Any,
        num_reasks: int = 1,
        api: Optional[Callable] = None,
        *args: Any,
        **kwargs: Any,
    ) -> GuardrailsOutputParser:
        try:
            from guardrails import Guard
        except ImportError:
            raise ImportError(
                "guardrails-ai package not installed. "
                "Install it by running `pip install guardrails-ai`."
            )
        return cls(
            guard=Guard.from_pydantic(output_class, "", num_reasks=num_reasks),
            api=api,
            args=args,
            kwargs=kwargs,
        )

    def get_format_instructions(self) -> str:
        return self.guard.raw_prompt.format_instructions

    def parse(self, text: str) -> Dict:
        return self.guard.parse(text, llm_api=self.api, *self.args, **self.kwargs)