Demo testing of anthropic's prompt caching for conversations

prompt_caching_conversation_demo.py

import asyncio
import aiohttp
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

API_KEY = os.getenv("ANTHROPIC_API_KEY")
API_URL = "https://api.anthropic.com/v1/messages"


async def make_request(session, system_prompt, messages):
    headers = {
        "Content-Type": "application/json",
        "X-API-Key": API_KEY,
        "anthropic-version": "2023-06-01",
        "anthropic-beta": "prompt-caching-2024-07-31"
    }

    payload = {
        "model": "claude-3-5-sonnet-20240620",
        "max_tokens": 1024,
        "system": [
            {
                "type": "text",
                "text": system_prompt,
                "cache_control": {"type": "ephemeral"}
            }
        ],
        "messages": messages
    }

    async with session.post(API_URL, headers=headers, json=payload) as response:
        result = await response.json()
        return result


def apply_sliding_window_cache_control(messages, window_size=3):
    for msg in messages:
        if "cache_control" in msg["content"][0]:
            del msg["content"][0]["cache_control"]

    for msg in messages[-window_size:]:
        msg["content"][0]["cache_control"] = {"type": "ephemeral"}

    return messages

async def main():
    system_prompt = """
    You are an AI assistant simulating a software engineering interview for TechInnovate, a leading technology company specializing in cloud-based solutions and artificial intelligence. The conversation will alternate between an interviewer and a candidate. Provide realistic and detailed responses for both roles.

    Company Information:
    TechInnovate is a rapidly growing tech company founded in 2010. With over 5,000 employees worldwide, it has offices in San Francisco, New York, London, and Singapore. The company's primary focus is on developing cutting-edge cloud infrastructure and AI-powered analytics tools for businesses of all sizes.

    Key products include:
    1. CloudScale: A highly scalable cloud computing platform
    2. AIInsight: An advanced analytics and machine learning tool
    3. SecureConnect: An enterprise-grade cybersecurity solution

    TechInnovate's tech stack includes:
    - Backend: Python, Go, Java
    - Frontend: React, Vue.js
    - Databases: PostgreSQL, MongoDB, Cassandra
    - Cloud infrastructure: AWS, Google Cloud Platform
    - AI/ML: TensorFlow, PyTorch, scikit-learn
    - DevOps: Docker, Kubernetes, Jenkins, GitLab CI

    Company Culture:
    TechInnovate prides itself on fostering a culture of innovation, collaboration, and continuous learning. The company offers:
    - Flexible work arrangements (remote-friendly)
    - Generous professional development budget
    - Regular hackathons and innovation challenges
    - Comprehensive health and wellness programs
    - Competitive salary and equity packages

    Job Role: Senior Software Engineer (Python Specialist)

    Responsibilities:
    1. Design, develop, and maintain scalable backend services using Python
    2. Collaborate with cross-functional teams to define and implement new features
    3. Optimize application performance and improve code quality
    4. Contribute to the architecture of large-scale distributed systems
    5. Mentor junior developers and participate in code reviews
    6. Stay updated with the latest trends in Python development and cloud technologies

    Requirements:
    1. 5+ years of experience in software development, with a strong focus on Python
    2. Extensive knowledge of Python frameworks (e.g., Django, Flask, FastAPI)
    3. Experience with cloud platforms (AWS, GCP) and containerization (Docker, Kubernetes)
    4. Solid understanding of database systems (SQL and NoSQL)
    5. Familiarity with CI/CD pipelines and DevOps practices
    6. Strong problem-solving skills and attention to detail
    7. Excellent communication and teamwork abilities
    8. BSc/MSc in Computer Science or a related field (or equivalent practical experience)

    Preferred Qualifications:
    1. Experience with microservices architecture
    2. Knowledge of machine learning libraries (e.g., TensorFlow, PyTorch)
    3. Contributions to open-source projects
    4. Experience with Go or Java

    The interviewer should ask questions relevant to this role and company, while the candidate should tailor their responses to demonstrate their fit for the position at TechInnovate.
    """

    conversation = [
        {"role": "user", "content": [{"type": "text", "text": "Interviewer: Welcome to TechInnovate! We're excited to discuss your candidacy for our Senior Software Engineer position. Can you start by telling me about your experience with Python and how it aligns with our tech stack?"}]},
        {"role": "assistant", "content": [{"type": "text", "text": "Candidate: Thank you for having me! I'm thrilled to be interviewing with TechInnovate. I have over 6 years of experience working extensively with Python, which aligns well with your tech stack. [Rest of the candidate's response...]"}]},
        {"role": "user", "content": [{"type": "text", "text": "Interviewer: That's an impressive background! Thank you for providing such a comprehensive overview. Regarding our current challenges, we're working on scaling our AIInsight product to handle larger datasets while maintaining low latency. Can you walk me through how you would approach optimizing a Python-based data processing pipeline for big data analytics?"}]},
    ]

    interview_questions = [
        "Excellent approach! Now, let's dive into your experience with cloud technologies. How have you used AWS or GCP in your previous projects, and how do you think this experience would apply to our CloudScale product?",
        "That's very relevant experience. TechInnovate places a strong emphasis on code quality and mentorship. Can you describe your approach to code reviews and how you've mentored junior developers in the past?",
        "Interesting perspective on mentorship. Let's switch gears a bit. Can you tell me about a particularly challenging technical problem you've solved recently? How did you approach it?",
        "That's a great example of problem-solving. Now, considering our focus on AI and machine learning at TechInnovate, how familiar are you with integrating ML models into production systems?",
        "Thank you for sharing your ML experience. Let's talk about system design. Can you walk me through how you would design a high-throughput, low-latency microservice architecture for our CloudScale product?",
        "Excellent design considerations. Now, let's discuss your experience with data structures and algorithms. Can you explain a situation where you had to optimize an algorithm for better performance?",
        "That's a solid example of algorithm optimization. Shifting gears, how do you stay updated with the latest trends and technologies in software engineering, particularly in the Python ecosystem?",
        "Great to hear about your commitment to continuous learning. Now, let's discuss a scenario. If you were tasked with migrating a monolithic application to a microservices architecture, what would be your approach and what challenges would you anticipate?",
        "Thank you for that comprehensive answer. As we wrap up, do you have any questions about TechInnovate, our culture, or the Senior Software Engineer role that I can answer for you?"
    ]

    async with aiohttp.ClientSession() as session:
        for i, question in enumerate(interview_questions):
            print(f"\nTurn {i + 1}:")

            # Apply sliding window cache control
            conversation = apply_sliding_window_cache_control(conversation)

            response = await make_request(session, system_prompt, conversation)

            print(f"Response: {response}")

            if 'error' in response:
                print(f"Error: {response['error']}")
                break

            assistant_reply = response.get('content', [{}])[0].get('text', '')
            if not assistant_reply:
                print("Error: Unexpected response structure")
                break

            print(f"Assistant: {assistant_reply}")

            # Print token usage information
            usage = response.get('usage', {})
            print(f"Input tokens: {usage.get('input_tokens', 'N/A')}")
            print(f"Output tokens: {usage.get('output_tokens', 'N/A')}")
            print(f"Cache read tokens: {usage.get('cache_read_input_tokens', 'N/A')}")
            print(f"Cache creation tokens: {usage.get('cache_creation_input_tokens', 'N/A')}")

            # Add assistant's reply to conversation
            conversation.append({"role": "assistant", "content": [{"type": "text", "text": assistant_reply}]})

            # Add next interviewer question to conversation
            conversation.append({"role": "user", "content": [{"type": "text", "text": f"Interviewer: {question}"}]})

            # Pause to allow for reading the output
            await asyncio.sleep(2)


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