o1 Style Thinking with Chain-of-Thought Reasoning using Mirascope

!pip install "mirascope[groq]" 
!pip install datetime

from typing import Literal

from mirascope.core import groq
from pydantic import BaseModel, Field


history: list[dict] = []


class COTResult(BaseModel):
    title: str = Field(..., desecription="The title of the step")
    content: str = Field(..., description="The output content of the step")
    next_action: Literal["continue", "final_answer"] = Field(
        ..., description="The next action to take"
    )

@groq.call("llama-3.3-70b-versatile", json_mode=True, response_model=COTResult)
def cot_step(prompt: str, step_number: int, previous_steps: str) -> str:
    return f"""
    You are an expert AI assistant that explains your reasoning step by step.
    For this step, provide a title that describes what you're doing, along with the content.
    Decide if you need another step or if you're ready to give the final answer.

    Guidelines:
    - Use AT MOST 5 steps to derive the answer.
    - Be aware of your limitations as an LLM and what you can and cannot do.
    - In your reasoning, include exploration of alternative answers.
    - Consider you may be wrong, and if you are wrong in your reasoning, where it would be.
    - Fully test all other possibilities.
    - YOU ARE ALLOWED TO BE WRONG. When you say you are re-examining
        - Actually re-examine, and use another approach to do so.
        - Do not just say you are re-examining.

    IMPORTANT: Do not use code blocks or programming examples in your reasoning. Explain your process in plain language.

    This is step number {step_number}.

    Question: {prompt}

    Previous steps:
    {previous_steps}
    """


@groq.call("llama-3.3-70b-versatile")
def final_answer(prompt: str, reasoning: str) -> str:
    return f"""
    Based on the following chain of reasoning, provide a final answer to the question.
    Only provide the text response without any titles or preambles.
    Retain any formatting as instructed by the original prompt, such as exact formatting for free response or multiple choice.

    Question: {prompt}

    Reasoning:
    {reasoning}

    Final Answer:
    """

def generate_cot_response(
    user_query: str,
) -> tuple[list[tuple[str, str, float]], float]:
    steps: list[tuple[str, str, float]] = []
    total_thinking_time: float = 0.0
    step_count: int = 1
    reasoning: str = ""
    previous_steps: str = ""

    while True:
        start_time: datetime = datetime.now()
        cot_result = cot_step(user_query, step_count, previous_steps)
        end_time: datetime = datetime.now()
        thinking_time: float = (end_time - start_time).total_seconds()

        steps.append(
            (
                f"Step {step_count}: {cot_result.title}",
                cot_result.content,
                thinking_time,
            )
        )
        total_thinking_time += thinking_time

        reasoning += f"\n{cot_result.content}\n"
        previous_steps += f"\n{cot_result.content}\n"

        if cot_result.next_action == "final_answer" or step_count >= 5:
            break

        step_count += 1

    # Generate final answer
    start_time = datetime.now()
    final_result: str = final_answer(user_query, reasoning).content
    end_time = datetime.now()
    thinking_time = (end_time - start_time).total_seconds()
    total_thinking_time += thinking_time

    steps.append(("Final Answer", final_result, thinking_time))

    return steps, total_thinking_time


def display_cot_response(
    steps: list[tuple[str, str, float]], total_thinking_time: float
) -> None:
    for title, content, thinking_time in steps:
        print(f"{title}:")
        print(content.strip())
        print(f"**Thinking time: {thinking_time:.2f} seconds**\n")

    print(f"**Total thinking time: {total_thinking_time:.2f} seconds**")

def run() -> None:
    question: str = "If a train leaves City A at 9:00 AM traveling at 60 km/h, and another train leaves City B (which is 300 km away from City A) at 10:00 AM traveling at 90 km/h toward City A, at what time will the trains meet?"
    print("(User):", question)
    # Generate COT response
    steps, total_thinking_time = generate_cot_response(question)
    display_cot_response(steps, total_thinking_time)

    # Add the interaction to the history
    history.append({"role": "user", "content": question})
    history.append(
        {"role": "assistant", "content": steps[-1][1]}
    )  # Add only the final answer to the history


# Run the function

run()