from __future__ import annotations

from dataclasses import dataclass

from frisk_sdk import FriskSession
from .frisk_callback_handler import FriskCallbackHandler


@dataclass
class FriskLangchainContext:
    frisk_session_id: str


class FriskLangchainSession(FriskSession):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.callbacks = FriskCallbackHandler(self.frisk, self)
        self.context = FriskLangchainContext(frisk_session_id=self.id)
