"""Chat with local and remote LLMs in the context of the wordslab-notebooks environment"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_chat.ipynb.

# %% auto 0
__all__ = ['NbChat']

# %% ../nbs/02_chat.ipynb 3
import os
from ollama import chat
import openai

from .core import WordslabNotebooks
from .context import get_notebook_context

from IPython.display import display, Markdown, clear_output

# %% ../nbs/02_chat.ipynb 49
class NbChat:
    def __init__(self, model=None):
        if model is None:
            wlnb = WordslabNotebooks()
            self.model = wlnb.default_model_chat
        else:
            self.model = model

        self.prompt_template = """You are an AI assistant designed to help the user learn and solve problems interactively.
You run in an interactive Jupyter notebook environment where the user can write code, take notes, and chat with you.
You work with the user step-by-step rather than just giving complete answers. You are especially good at:
- Breaking down complex topics into manageable pieces
- Helping the user work through coding problems in Python
- Encouraging the user to try things himself, with guidance when he needs it
- Adapting to the user level and interests
You are designed to be collaborative - you ask questions, check the user understanding, and let him explore ideas rather than just lecturing. 
You can help with teaching, coding, problem-solving, research, and creative projects.
What sets you apart is your teaching approach - you focus on helping the user develop his skills rather than just giving him answers. 
You provide information in small chunks, check in frequently to see if things make sense, and encourage the user to try things himself.

# Jupyter notebook - all cells above the user instruction in XML

{notebook_context}

# User Instruction - in the last code cell of the notebook

Execute this user instruction in the context of the code cells above:

{user_instruction}
"""

    async def __call__(self, user_instruction, timeout=1):
        notebook_context = await get_notebook_context(timeout=timeout)
        prompt = self.prompt_template.format(user_instruction=user_instruction, notebook_context=notebook_context)
        stream = chat(model=self.model, messages=[{'role': 'user', 'content': prompt}], stream=True)
        streamed_text = ""
        for chunk in stream:
            streamed_text += chunk['message']['content']    
            clear_output(wait=True)
            display(Markdown(streamed_text))
