import time
import panel as pn

pn.extension()


def callback(contents, user, instance):
    time.sleep(0.5)

    card = pn.Card(
        title="Thinking...", sizing_mode="stretch_width",
    )

    accordion = pn.Accordion(sizing_mode="stretch_width", active=[0], margin=0)
    card.append(accordion)
    yield card

    first_card = pn.pane.Markdown("Thinking...")
    accordion.append(("What are my meetings today?", first_card))
    time.sleep(2)
    first_card.object = "You have a meeting at 10:00 AM with the team and another at 2:00 PM with the client."
    
    second_card = pn.pane.Markdown("Thinking...")
    accordion.append(("What time is it right now", second_card))
    time.sleep(1)
    second_card.object = "It is currently 9:30 AM."

    third_card = pn.pane.Markdown("Which meeting is next?")
    accordion.append(("Which meeting is next?", third_card))
    time.sleep(1)
    third_card.object = "Your next meeting is at 10:00 AM with the team."

    fourth_card = pn.pane.Markdown("Thinking...")
    accordion.append(("How long until the next meeting?", fourth_card))
    time.sleep(1)
    fourth_card.object = "Your next meeting is in 30 minutes."

    return "Your next meeting is in 30 minutes."

chat = pn.chat.ChatInterface(callback=callback)
chat.show()
chat.send("How long until my next meeting?")

import time
import panel as pn
from contextlib import contextmanager

pn.extension()


def append_step(accordion, question, response, placeholder="Thinking..."):
    card = pn.pane.Markdown(placeholder)
    accordion.append((question, card))
    time.sleep(1)
    accordion.active = [len(accordion.objects) - 1]  # configurable
    time.sleep(1)
    for i, char in enumerate(response):
        if i == 0:
            card.object = char
        else:
            time.sleep(0.001)
            card.object += char


@contextmanager
def show_steps():
    try:
        card = pn.Card(
            title="Thinking...",
            sizing_mode="stretch_width",
        )
        accordion = pn.Accordion(sizing_mode="stretch_width", margin=0, toggle=True)
        card.append(accordion)
        yield card, accordion
    finally:
        card.title = accordion.objects[-1].object


def callback(contents, user, instance):
    time.sleep(0.5)
    with show_steps() as (card, accordion):
        yield accordion

        # in actual class it would be
        # accordion.append(markdown_obj, markdown_obj)
        append_step(
            accordion,
            "What are my meetings today?",
            "You have a meeting at 10:00 AM with the team and another at 2:00 PM with the client.",
        )

        append_step(
            accordion,
            "What time is it right now",
            "It is currently 9:30 AM.",
        )

        append_step(
            accordion,
            "Which meeting is next?",
            "Your next meeting is at 10:00 AM with the team.",
        )

        append_step(
            accordion,
            "How long until the next meeting?",
            "Your next meeting is in 30 minutes.",
        )


chat = pn.chat.ChatInterface(callback=callback, callback_exception="verbose")
chat.show()

from typing import List
from pydantic import Field, BaseModel


class Query(BaseModel):
    """Class representing a single question in a query plan."""

    id: int = Field(..., description="Unique id of the query")
    question: str = Field(
        ...,
        description="Question asked using a question answering system",
    )
    dependencies: List[int] = Field(
        default_factory=list,
        description="List of sub questions that need to be answered before asking this question",
    )


class QueryPlan(BaseModel):
    """Container class representing a tree of questions to ask a question answering system."""

    query_graph: List[Query] = Field(
        ..., description="The query graph representing the plan"
    )

    def _dependencies(self, ids: List[int]) -> List[Query]:
        """Returns the dependencies of a query given their ids."""
        return [q for q in self.query_graph if q.id in ids]


import instructor
from openai import OpenAI

# Apply the patch to the OpenAI client
# enables response_model keyword
client = instructor.patch(OpenAI())


def query_planner(question: str) -> QueryPlan:
    PLANNING_MODEL = "gpt-4-0613"

    messages = [
        {
            "role": "system",
            "content": "You are a world class query planning algorithm capable ofbreaking apart questions into its dependency queries such that the answers can be used to inform the parent question. Do not answer the questions, simply provide a correct compute graph with good specific questions to ask and relevant dependencies. Before you call the function, think step-by-step to get a better understanding of the problem.",
        },
        {
            "role": "user",
            "content": f"Consider: {question}\nGenerate the correct query plan.",
        },
    ]

    root = client.chat.completions.create(
        model=PLANNING_MODEL,
        temperature=0,
        response_model=instructor.Partial[QueryPlan],
        messages=messages,
        max_tokens=1000,
        stream=True,
    )
    return root


import time
import panel as pn
from contextlib import contextmanager

pn.extension()


@contextmanager
def show_steps(instance):
    try:
        card = pn.Card(
            title="Thinking...",
            sizing_mode="stretch_width",
        )
        instance.stream(card)
        yield card
    finally:
        card.title = card.objects[-1].objects[-1].object


def callback(contents, user, instance):

    queries = {}
    with show_steps(instance) as card:
        chunked_plan = query_planner(
            "What is the difference in populations of Canada and the Andrew's home country?"
        )
        for chunk in chunked_plan:
            for query in chunk.model_dump()["query_graph"] or []:
                if query.get("id") is None or hasattr(query, "question"):
                    continue

                if query["id"] not in queries:
                    answer_md = pn.pane.Markdown("")
                    sub_card = pn.Card(
                        answer_md,
                        title=query["question"] or "",
                        sizing_mode="stretch_width",
                    )
                    queries[query["id"]] = sub_card
                    card.append(sub_card)
                else:
                    sub_card = queries[query["id"]]
                    sub_card.title = query["question"] or ""

        history = []
        for sub_card in card:
            # serialize is desperately needed here:
            question_dict = {"role": "user", "content": sub_card.title}
            history.append(question_dict)
            messages = [
                {
                    "role": "system",
                    "content": "Answer the following question to the best of your abilities.",
                },
                {
                    "role": "user",
                    "content": "Andrew lives in the US"
                },
                *history,
            ]
            response = client.chat.completions.create(
                model="gpt-3.5-turbo",
                temperature=0,
                messages=messages,
                max_tokens=1000,
                stream=True,
            )
            for chunk in response:
                if chunk:
                    sub_card.objects[0].object += chunk.choices[0].delta.content or ""
            history.append({"role": "assistant", "content": sub_card.objects[0].object})


chat = pn.chat.ChatInterface(callback=callback, callback_exception="verbose")
chat.show()