What is Retrieval Augmented Thinking (RAT) and how does it work?

import os
from litellm import completion
from openai import OpenAI
import weave
import time
from transformers import AutoTokenizer

import weave; weave.init("rat")


# Model Constants
DEEPSEEK_MODEL = "deepseek-reasoner"
GPT_MODEL = "gpt-4o-mini"

# Initialize DeepSeek client and tokenizer
weave.init("deepseek_r1_api_streaming")
r1_api_key = "your deepseek api key"  # Replace with your API key
deepseek_client = OpenAI(api_key=r1_api_key, base_url="https://api.deepseek.com")

# Initialize DeepSeek tokenizer
tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-V2", trust_remote_code=True)


@weave.op
def get_deepseek_reasoning(prompt, max_think_tokens=1000, max_think_time=60):

    start_time = time.time()
    
    print("\nReasoning Process:")
    
    modified_prompt = f"{prompt}\n\nPlease limit your response to approximately {max_think_tokens} tokens."
    
    # Use direct OpenAI client for DeepSeek
    response = deepseek_client.chat.completions.create(
        model=DEEPSEEK_MODEL,
        max_tokens=1,  # This triggers reasoner mode
        messages=[{"role": "user", "content": modified_prompt}],
        stream=True
    )

    reasoning_content = ""
    final_content = ""
    token_count = 0
    limit_reached = False
    limit_reason = ""

    try:
        # Process chunks with timeout and token limit
        for chunk in response:
            # Check if we've exceeded the max thinking time
            current_time = time.time()
            elapsed = current_time - start_time
            
            if elapsed > max_think_time and not limit_reached:
                limit_reached = True
                limit_reason = "TIME"
                print("\n\n[TIME LIMIT REACHED]")
                # Explicitly break out of the loop on time limit
                break
                
            # Process the chunk
            if chunk.choices[0].delta.reasoning_content:
                reasoning_piece = chunk.choices[0].delta.reasoning_content
                
                # Count tokens in the new piece and check token limit
                new_tokens = len(tokenizer.encode(reasoning_piece))
                token_count += new_tokens
                
                if token_count > max_think_tokens and not limit_reached:
                    limit_reached = True
                    limit_reason = "TOKEN"
                    print("\n\n[TOKEN LIMIT REACHED]")
                    # Explicitly break out of the loop on token limit
                    break
                    
                reasoning_content += reasoning_piece
                print(reasoning_piece, end="", flush=True)
            elif chunk.choices[0].delta.content:
                final_content += chunk.choices[0].delta.content
    
    finally:
        # Properly terminate the stream
        try:
            # This closes the HTTP connection for the stream
            response.close()
            print("[Stream terminated]")
        except Exception as e:
            print(f"Error closing stream: {e}")

    # Calculate and display elapsed time
    elapsed_time = time.time() - start_time
    if elapsed_time >= 60:
        time_str = f"{elapsed_time/60:.1f} minutes"
    else:
        time_str = f"{elapsed_time:.1f} seconds"
    
    print(f"\n\nThought for {time_str} ({token_count} tokens)")
    if limit_reached:
        print(f"Stopped due to {limit_reason} limit")
    print("\n")
    
    return reasoning_content


@weave.op
def get_gpt_response(prompt, reasoning):
    """Get response from GPT model via LiteLLM"""
    combined_prompt = (
        f"<question>{prompt}</question>\n\n"
        f"<thinking>{reasoning}</thinking>\n\n"
    )
    
    print(f"\n{GPT_MODEL}:")
    
    try:
        # Using LiteLLM for GPT
        completion_response = completion(
            model=GPT_MODEL,
            messages=[{"role": "user", "content": combined_prompt}],
            stream=True
        )
        
        full_response = ""
        for chunk in completion_response:
            try:
                if hasattr(chunk, 'choices') and len(chunk.choices) > 0:
                    delta = chunk.choices[0].delta
                    if hasattr(delta, 'content') and delta.content:
                        content_piece = delta.content
                        full_response += content_piece
                        print(content_piece, end="", flush=True)
            except Exception as e:
                print(f"\nError processing chunk: {str(e)}")
                continue
                
    except Exception as e:
        print(f"\nError in streaming response: {str(e)}")
        return "Error occurred while streaming response"
    
    print("\n")
    return full_response


@weave.op
def get_deepseek_final_answer(prompt, reasoning):
    """
    Get a final answer from DeepSeek using the reasoning
    """
    combined_prompt = (
        f"<question>{prompt}</question>\n\n"
        f"<thinking>{reasoning}</thinking>\n\n"
        f"Based on the above reasoning, provide a clear and concise answer."
    )
    
    print(f"\nDeepSeek Final Answer:")
    
    try:
        # Use DeepSeek for the final answer
        response = deepseek_client.chat.completions.create(
            model="deepseek-chat",  # Using the chat model for the final answer
            messages=[{"role": "user", "content": combined_prompt}],
            temperature=0.7,
            max_tokens=500,
            stream=True
        )
        
        full_response = ""
        for chunk in response:
            try:
                if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
                    content_piece = chunk.choices[0].delta.content
                    full_response += content_piece
                    print(content_piece, end="", flush=True)
            except Exception as e:
                print(f"\nError processing chunk: {str(e)}")
                continue
                
    except Exception as e:
        print(f"\nError in streaming response: {str(e)}")
        return "Error occurred while streaming response"
    
    print("\n")
    return full_response

def main():
    # Hardcoded prompt
    prompt = "Explain the concept of quantum computing in simple terms."
    
    # Process with both token and time limits (whichever comes first)
    reasoning = get_deepseek_reasoning(
        prompt, 
        max_think_tokens=500,  # Limit to 500 tokens
        max_think_time=500      # Limit to 30 seconds
    )
    
    # Choose either GPT or DeepSeek for the final answer
    use_gpt = False
    
    if use_gpt:
        # Option 1: Use GPT for the final answer
        final_answer = get_gpt_response(prompt, reasoning)
    else:
        # Option 2: Use DeepSeek for the final answer
        final_answer = get_deepseek_final_answer(prompt, reasoning)

if __name__ == "__main__":
    main()

import asyncio
import time
from anthropic import Anthropic
from openai import OpenAI
import weave
import json
from datasets import load_dataset

# Initialize Weave
weave.init("aime_evaluation")

# Claude 3.7 configuration
CLAUDE_API_KEY = "your claude api key"  # Replace with your actual API key
CLAUDE_MODEL = "claude-3-7-sonnet-20250219"

THINKING_BUDGET = 4000
MAX_TOKENS = THINKING_BUDGET + 4000


# Initialize Anthropic client for Claude 3.7
claude_client = Anthropic(api_key=CLAUDE_API_KEY)

# Initialize OpenAI client for O3-mini (mistake detection)
openai_client = OpenAI(
    api_key="your OpenAI api key"  # Replace with your actual API key
)

# # Initialize OpenAI client for GPT-4o (scoring)
# gpt4o_client = OpenAI(
#     api_key="your_gpt4o_api_key",  # Replace with your GPT-4o API key
# )

# System prompt for Claude
SYSTEM_MESSAGE = "Solve the following problem. Put your final answer within \\boxed{}: "

class ClaudeO3Verifier(weave.Model):
    """
    A simplified model that combines Claude 3.7 for reasoning with O3-mini for verification
    and error correction, following the "solve-critique-fix" loop.
    """
    
    @staticmethod
    async def get_claude_reasoning(prompt: str, critic=None, verbose=False) -> dict:
        """
        Get reasoning from Claude 3.7 with optional critic feedback.
        """
        try:
            start_time = time.time()
            
            # Prepare full prompt with critic feedback if available
            full_prompt = prompt
            if critic and critic != "Correct Reasoning" and critic != "":
                full_prompt = f"""Problem: {prompt}

A review of your previous solution found this error:
{critic}

Please solve this problem again with correct reasoning."""
            
            # Using streaming mode with Claude 3.7
            full_response = ""
            thinking_content = ""
            
            if verbose:
                print(f"\nStreaming Claude 3.7 Reasoning for problem: {prompt[:50]}...")
                print("-" * 50)
            
            # Stream Claude's response
            with claude_client.messages.stream(
                model=CLAUDE_MODEL,
                max_tokens=MAX_TOKENS,
                thinking={"type": "enabled", "budget_tokens": THINKING_BUDGET},
                messages=[{"role": "user", "content": SYSTEM_MESSAGE + full_prompt}]
            ) as stream:
                current_block_type = None
                for event in stream:
                    if event.type == "content_block_start":
                        current_block_type = event.content_block.type
                        if verbose:
                            print(f"\n--- Starting {current_block_type} block ---")
                    elif event.type == "content_block_delta":
                        if event.delta.type == "thinking_delta":
                            thinking_content += event.delta.thinking
                            if verbose and len(thinking_content) % 100 == 0:
                                print(".", end="", flush=True)
                        elif event.delta.type == "text_delta":
                            full_response += event.delta.text
                            if verbose and len(full_response) % 100 == 0:
                                print("*", end="", flush=True)
                    elif event.type == "content_block_stop":
                        if verbose:
                            print(f"\n--- End {current_block_type} block ---")
                    elif event.type == "message_stop":
                        if verbose:
                            print("\n--- Message complete ---")
            
            end_time = time.time()
            if verbose:
                print("\n" + "-" * 50)
                print(f"Total time: {end_time - start_time:.2f} seconds")
                print(f"Total thinking length: {len(thinking_content)} characters")
                print(f"Total response length: {len(full_response)} characters")
            
            return {
                "thinking": thinking_content,
                "response": full_response
            }
            
        except Exception as e:
            end_time = time.time()
            print(f"\nError after {end_time - start_time:.2f} seconds: {str(e)}")
            return {"thinking": "", "response": f"Error: {str(e)}"}
    @weave.op
    @staticmethod
    async def check_reasoning_correctness(problem: str, reasoning: dict) -> str:
        """
        Check the full reasoning for correctness using O3-mini.
        Returns "Correct Reasoning" or "Error - {description}".
        """
        if not reasoning or (not reasoning.get("thinking") and not reasoning.get("response")):
            return "Error - Empty reasoning provided"
        
        # Prioritize thinking content, but fall back to response if thinking is empty
        reasoning_text = reasoning.get("thinking") or reasoning.get("response")
        
        # Clean reasoning if needed
        clean_reasoning = reasoning_text.strip()
        
        validation_prompt = "\n".join([
            "Examine this complete mathematical reasoning trace for any logical mistakes, calculation errors, or incorrect assumptions:",
            "",
            "the problem: ", 
            problem, 
            "reasoning: ",
            clean_reasoning,
            "",
            "Focus on mathematical correctness. Pay special attention to:",
            "1. Arithmetic errors",
            "2. Algebraic manipulation errors",
            "3. Logical flaws in the reasoning",
            "4. Incorrect application of formulas",
            "5. Unwarranted assumptions",
            "",
            "Respond with EXACTLY:",
            "- \"Correct Reasoning\" if no mistakes are found",
            "- \"Error - [a explanation of which step the error occurs, detailed explanation of the specific error, along with an explanation of what to do instead]\" if you find any mistakes"
        ])


        try:
            response = openai_client.chat.completions.create(
                model="o3-mini",
                messages=[{"role": "user", "content": validation_prompt}]
            )
            result = response.choices[0].message.content.strip()
            
            if result.lower().startswith("correct reasoning"):
                return "Correct Reasoning"
            elif result.lower().startswith("error"):
                return result
            else:
                # Force result into expected format
                if "error" in result.lower():
                    return f"Error - DO NOT REPEAT THIS MISTAKE AGAIN - {result}"
                else:
                    return "Correct Reasoning"  # Default to correct if unclear
                    
        except Exception as e:
            print(f"Failed to validate reasoning: {e}")
            return f"Error - Validation failed: {str(e)}"

    @weave.op
    async def predict(self, text: str) -> str:
        """
        Run the complete inference pipeline using iterative solve-critique-fix loop.
        
        Args:
            text: The mathematical problem to solve
            
        Returns:
            str: The final solution to the problem
        """
        try:
            print(f"Problem: {text[:100]}...")
            
            # Initialize loop variables
            max_iterations = 5  # Maximum number of attempts
            current_iteration = 0
            critic = ""
            done = False
            final_reasoning = None
            partial_critic = None 
            
            # Iterative solve-critique-fix loop
            while not done and current_iteration < max_iterations:
                current_iteration += 1
                print(f"\nIteration {current_iteration}/{max_iterations}")
                
                # Get reasoning from Claude 3.7
                print("Getting reasoning from Claude 3.7...")
                reasoning = await self.get_claude_reasoning(text, critic, verbose=False)
                final_reasoning = reasoning
                
                if not reasoning.get("response"):
                    return "Failed to get reasoning"
                
                # Check reasoning with O3-mini
                print("Checking reasoning with O3-mini...")
                partial_critic = await ClaudeO3Verifier.check_reasoning_correctness(text, reasoning)
                
                print(f"Critic: {critic[:100]}...")
                
                # Check if we're done
                if partial_critic == "Correct Reasoning" or current_iteration == max_iterations:
                    done = True
            
                critic += "Mistake: " + partial_critic
            # Return the final response for evaluation
            return final_reasoning.get("response", "No response generated")
            
        except Exception as e:
            print(f"Critical error in predict method: {str(e)}")
            return f"Error processing problem. Details: {str(e)}"

@weave.op
async def gpt4o_scorer(label: str, model_output: str) -> dict:
    """Score the model's output by comparing it with the ground truth."""
    query = f"""
    YOU ARE A LLM JUDGE DETERMINING IF THE FOLLOWING MODEL GENERATED ANSWER IS THE SAME AS THE CORRECT ANSWER 
    I WILL GIVE YOU THE LAST 100 CHARS OF THE MODEL'S REASONING PATH, WHICH WILL CONTAIN THE FINAL ANSWER -> 

    Model's Answer (last 100 chars): {str(model_output)[-100:]}
    Correct Answer: {label}
    
    Your task:
    1. State the model's predicted answer (answer only).
    2. State the ground truth (answer only).
    3. Determine if the model's final answer is correct (ignore formatting differences, etc.). RESPOND with the predicted and ground truth answer, followed with a JSON object containing the correctness encapsulated within the following delimiters:
       ```json
       {{ "correctness": true/false }}
       ```
    """
    
    try:
        # Perform inference using GPT-4o client
        response = openai_client.chat.completions.create(
            model="gpt-4o-2024-08-06",
            messages=[{"role": "user", "content": query}]
        )
        
        result = response.choices[0].message.content
        
        # Extract correctness JSON object from the response
        try:
            json_start = result.index("```json") + 7
            json_end = result.index("```", json_start)
            correctness = json.loads(result[json_start:json_end].strip()).get("correctness", False)
        except (ValueError, IndexError, json.JSONDecodeError):
            correctness = False
            print("Failed to parse correctness JSON")

        return {"correctness": correctness, "reasoning": result}
    
    except Exception as e:
        print(f"Scoring error: {e}")
        return {"correctness": False, "reasoning": f"Scoring error: {str(e)}"}

# Load and preprocess dataset
def load_ds():
    print("Loading AIME dataset...")
    dataset = load_dataset("Maxwell-Jia/AIME_2024")["train"]  # no test set here 
    return [{"text": row["Problem"], "label": row["Answer"]} for row in dataset]

async def run_evaluations():
    print("Loading dataset...")
    dataset = load_ds()
    
    print(f"Dataset loaded with {len(dataset)} problems")
    
    # For testing purposes, you might want to use a smaller subset initially
    # Uncomment the next line to use a smaller subset
    # dataset = dataset[:5]  # Use only 5 problems for testing
    
    print("Initializing model...")
    model = ClaudeO3Verifier()
    
    print("Preparing dataset for evaluation...")
    dataset_prepared = [{"text": row["text"], "label": row["label"]} for row in dataset]
    
    print("Running evaluation...")
    scorers = [gpt4o_scorer]
    
    evaluation = weave.Evaluation(
        dataset=dataset_prepared,
        scorers=scorers,
        name="Claude 3.7 AIME Evaluation"
    )
    
    print("Starting evaluation...")
    results = await evaluation.evaluate(model)
    print(f"Evaluation complete!")
    print(f"Results: {results}")
    
    return results

if __name__ == "__main__":
    asyncio.run(run_evaluations())