from literalai import LiteralClient
import os
from .base import ChatProcessorBase


class LiteralaiChatProcessor(ChatProcessorBase):
    def __init__(self, tags=None):
        self.literal_client = LiteralClient(api_key=os.getenv("LITERAL_API_KEY"))
        self.literal_client.reset_context()
        with self.literal_client.thread(name="TEST") as thread:
            self.thread_id = thread.id
            self.thread = thread
            if tags is not None and type(tags) == list:
                self.thread.tags = tags
        print(f"Thread ID: {self.thread}")

    def process(self, user_message, assistant_message, source_dict):
        with self.literal_client.thread(thread_id=self.thread_id) as thread:
            self.literal_client.message(
                content=user_message,
                type="user_message",
                name="User",
            )
            self.literal_client.message(
                content=assistant_message,
                type="assistant_message",
                name="AI_Tutor",
            )

    async def rag(self, user_query: str, chain, cb):
        with self.literal_client.step(
            type="retrieval", name="RAG", thread_id=self.thread_id
        ) as step:
            step.input = {"question": user_query}
            res = await chain.acall(user_query, callbacks=[cb])
            step.output = res
        return res