# Code generated by Lark OpenAPI.

from typing import Any, Optional, Union, Dict, List, Set, IO, Callable, Type
from lark_oapi.core.construct import init
from .passage import Passage
from .rag_probe import RagProbe


class RagAnswerResponse(object):
    _types = {
        "answer": str,
        "passages": List[Passage],
        "probe": RagProbe,
    }

    def __init__(self, d=None):
        self.answer: Optional[str] = None
        self.passages: Optional[List[Passage]] = None
        self.probe: Optional[RagProbe] = None
        init(self, d, self._types)

    @staticmethod
    def builder() -> "RagAnswerResponseBuilder":
        return RagAnswerResponseBuilder()


class RagAnswerResponseBuilder(object):
    def __init__(self) -> None:
        self._rag_answer_response = RagAnswerResponse()

    def answer(self, answer: str) -> "RagAnswerResponseBuilder":
        self._rag_answer_response.answer = answer
        return self

    def passages(self, passages: List[Passage]) -> "RagAnswerResponseBuilder":
        self._rag_answer_response.passages = passages
        return self

    def probe(self, probe: RagProbe) -> "RagAnswerResponseBuilder":
        self._rag_answer_response.probe = probe
        return self

    def build(self) -> "RagAnswerResponse":
        return self._rag_answer_response
