Prompt Engineering

"""
Anthropic Claude-Specific Prompt Engineering Examples

Demonstrates Claude's unique features:
- System prompts
- XML tag patterns
- Tool use / function calling
- Streaming responses
- Extended thinking
- Prompt caching

Installation:
    pip install anthropic

Usage:
    export ANTHROPIC_API_KEY="your-api-key"
    python anthropic-examples.py
"""

import os
import anthropic
from typing import List, Dict, Any


# Initialize client
client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))


# ============================================================================
# Example 1: Basic Claude Message
# ============================================================================

def basic_message_example():
    """Simple message with system prompt."""
    print("\n=== Basic Message Example ===")

    message = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=1024,
        system="You are a helpful AI assistant specialized in Python programming.",
        messages=[
            {"role": "user", "content": "Write a function to calculate factorial"}
        ]
    )

    print(message.content[0].text)


# ============================================================================
# Example 2: XML-Structured Prompts (Claude Best Practice)
# ============================================================================

def xml_structured_prompt():
    """Use XML tags for better structure (Claude's strength)."""
    print("\n=== XML-Structured Prompt ===")

    system_prompt = """
You are a document analyzer. Follow these guidelines:

<capabilities>
- Extract key information from documents
- Summarize content concisely
- Identify main themes and topics
</capabilities>

<output_format>
Provide analysis in this structure:
<analysis>
  <summary>Brief summary</summary>
  <key_points>
    <point>Key point 1</point>
    <point>Key point 2</point>
  </key_points>
  <themes>List of themes</themes>
</analysis>
</output_format>
"""

    user_message = """
<document>
<title>The Future of AI</title>
<content>
Artificial intelligence is transforming industries worldwide. From healthcare
to finance, AI systems are becoming increasingly sophisticated. However,
challenges remain in ethics, regulation, and ensuring AI benefits all of society.
</content>
</document>

Analyze this document.
"""

    message = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=1024,
        system=system_prompt,
        messages=[{"role": "user", "content": user_message}]
    )

    print(message.content[0].text)


# ============================================================================
# Example 3: Tool Use / Function Calling
# ============================================================================

def tool_use_example():
    """Demonstrate Claude's tool use capabilities."""
    print("\n=== Tool Use Example ===")

    # Define tools
    tools = [
        {
            "name": "get_weather",
            "description": "Get current weather for a specific location. Use this when users ask about weather conditions.",
            "input_schema": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "City name, e.g., 'San Francisco, CA'"
                    },
                    "units": {
                        "type": "string",
                        "enum": ["celsius", "fahrenheit"],
                        "description": "Temperature unit"
                    }
                },
                "required": ["location"]
            }
        },
        {
            "name": "calculate",
            "description": "Perform mathematical calculations. Use for any arithmetic operations.",
            "input_schema": {
                "type": "object",
                "properties": {
                    "expression": {
                        "type": "string",
                        "description": "Mathematical expression to evaluate, e.g., '2 + 2 * 3'"
                    }
                },
                "required": ["expression"]
            }
        }
    ]

    # Make initial request
    message = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=1024,
        tools=tools,
        messages=[
            {"role": "user", "content": "What's the weather in Tokyo and what's 15 * 23?"}
        ]
    )

    print("Claude's response:")
    print(f"Stop reason: {message.stop_reason}")

    # Check if Claude wants to use tools
    if message.stop_reason == "tool_use":
        for block in message.content:
            if block.type == "tool_use":
                print(f"\nTool called: {block.name}")
                print(f"Tool input: {block.input}")

                # Simulate tool execution
                if block.name == "get_weather":
                    tool_result = {
                        "temperature": 18,
                        "condition": "Partly cloudy",
                        "humidity": 65
                    }
                elif block.name == "calculate":
                    import ast
                    tool_result = {"result": eval(block.input["expression"])}

                print(f"Tool result: {tool_result}")


# ============================================================================
# Example 4: Multi-Turn Conversation
# ============================================================================

def multi_turn_conversation():
    """Demonstrate conversation history management."""
    print("\n=== Multi-Turn Conversation ===")

    conversation_history = []

    # Turn 1
    user_message_1 = "My name is Alice and I love Python programming."
    conversation_history.append({"role": "user", "content": user_message_1})

    response_1 = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=256,
        messages=conversation_history
    )

    assistant_message_1 = response_1.content[0].text
    conversation_history.append({"role": "assistant", "content": assistant_message_1})

    print(f"User: {user_message_1}")
    print(f"Claude: {assistant_message_1}")

    # Turn 2
    user_message_2 = "What's my name and what language do I like?"
    conversation_history.append({"role": "user", "content": user_message_2})

    response_2 = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=256,
        messages=conversation_history
    )

    assistant_message_2 = response_2.content[0].text

    print(f"\nUser: {user_message_2}")
    print(f"Claude: {assistant_message_2}")


# ============================================================================
# Example 5: Streaming Responses
# ============================================================================

def streaming_example():
    """Stream Claude's response token-by-token."""
    print("\n=== Streaming Example ===")
    print("Claude's response (streaming): ", end="", flush=True)

    with client.messages.stream(
        model="claude-3-5-sonnet-20241022",
        max_tokens=512,
        messages=[
            {"role": "user", "content": "Write a haiku about artificial intelligence"}
        ]
    ) as stream:
        for text in stream.text_stream:
            print(text, end="", flush=True)

    print("\n")


# ============================================================================
# Example 6: Prompt Caching (Cost Optimization)
# ============================================================================

def prompt_caching_example():
    """Use prompt caching to reduce costs for repeated context."""
    print("\n=== Prompt Caching Example ===")

    # Large context that will be reused
    large_codebase = """
# Authentication Module
class UserAuth:
    def __init__(self, db):
        self.db = db

    def login(self, username, password):
        user = self.db.find_user(username)
        if user and user.verify_password(password):
            return self.create_session(user)
        return None

    def create_session(self, user):
        # Creates JWT token
        pass

# Payment Module
class PaymentProcessor:
    def process_payment(self, amount, card):
        # Process payment logic
        pass
"""

    # First call - cache is created
    print("First call (creates cache):")
    message1 = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=1024,
        system=[
            {
                "type": "text",
                "text": "You are a code reviewer specializing in Python."
            },
            {
                "type": "text",
                "text": f"Codebase to review:\n\n{large_codebase}",
                "cache_control": {"type": "ephemeral"}  # Cache this block
            }
        ],
        messages=[
            {"role": "user", "content": "Review the UserAuth class for security issues"}
        ]
    )

    print(f"Usage: {message1.usage}")
    print(f"Response: {message1.content[0].text[:200]}...\n")

    # Second call - uses cache (90% cost reduction on cached portion)
    print("Second call (uses cache):")
    message2 = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=1024,
        system=[
            {
                "type": "text",
                "text": "You are a code reviewer specializing in Python."
            },
            {
                "type": "text",
                "text": f"Codebase to review:\n\n{large_codebase}",
                "cache_control": {"type": "ephemeral"}  # Same cached block
            }
        ],
        messages=[
            {"role": "user", "content": "Review the PaymentProcessor class"}
        ]
    )

    print(f"Usage: {message2.usage}")
    print(f"Cache hits: {getattr(message2.usage, 'cache_read_input_tokens', 0)} tokens")


# ============================================================================
# Example 7: Extended Thinking (Claude's Chain-of-Thought)
# ============================================================================

def extended_thinking_example():
    """Use Claude's extended thinking for complex reasoning."""
    print("\n=== Extended Thinking Example ===")

    # Note: Extended thinking is available on Claude 3.5 Sonnet and Opus
    system_prompt = """
You are a mathematics tutor. When solving problems:
1. Show all your work step-by-step
2. Explain your reasoning at each step
3. Verify your answer

Use <thinking> tags to show your internal reasoning process.
"""

    message = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=2048,
        system=system_prompt,
        messages=[
            {
                "role": "user",
                "content": """
A train leaves Station A at 2:00 PM traveling at 60 mph.
Another train leaves Station B at 3:00 PM traveling at 80 mph toward Station A.
The stations are 300 miles apart.
At what time do the trains meet?
"""
            }
        ]
    )

    print(message.content[0].text)


# ============================================================================
# Example 8: Structured Data Extraction with Tool Use
# ============================================================================

def structured_extraction_example():
    """Extract structured data using tool use (Claude's JSON mode)."""
    print("\n=== Structured Data Extraction ===")

    # Define schema as a tool
    tools = [{
        "name": "record_user_info",
        "description": "Record structured user information",
        "input_schema": {
            "type": "object",
            "properties": {
                "name": {"type": "string", "description": "Full name"},
                "age": {"type": "integer", "description": "Age in years"},
                "email": {"type": "string", "description": "Email address"},
                "interests": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "List of interests"
                }
            },
            "required": ["name", "age", "email"]
        }
    }]

    message = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=1024,
        tools=tools,
        messages=[{
            "role": "user",
            "content": """
Extract information from this bio:

"Hi, I'm Sarah Johnson, a 28-year-old software engineer.
You can reach me at sarah.j@email.com.
I'm passionate about machine learning, rock climbing, and photography."
"""
        }]
    )

    # Extract structured data
    for block in message.content:
        if block.type == "tool_use":
            print(f"Extracted data:")
            import json
            print(json.dumps(block.input, indent=2))


# ============================================================================
# Example 9: RAG Pattern with Citations
# ============================================================================

def rag_with_citations():
    """Retrieval-augmented generation with source citations."""
    print("\n=== RAG with Citations ===")

    system_prompt = """
Answer questions using ONLY information from the provided documents.
Cite sources using <citation source="document_name">fact</citation> tags.
If information is not in the documents, say "I don't have this information."
"""

    user_message = """
<documents>
  <document id="1" source="product_manual.pdf">
    The XYZ-2000 operates at 120V AC and consumes 500W maximum power.
    Recommended for indoor use only.
  </document>
  <document id="2" source="safety_guide.pdf">
    Always unplug the device before performing any maintenance.
    Keep away from water and moisture.
  </document>
  <document id="3" source="warranty.pdf">
    Product includes a 2-year manufacturer warranty covering defects.
    Warranty void if device is opened by unauthorized personnel.
  </document>
</documents>

<question>
What voltage does the XYZ-2000 use, what safety precautions should I take,
and what's the warranty period?
</question>
"""

    message = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=1024,
        system=system_prompt,
        messages=[{"role": "user", "content": user_message}]
    )

    print(message.content[0].text)


# ============================================================================
# Example 10: Prefill Pattern (Control Output Format)
# ============================================================================

def prefill_pattern():
    """Use prefill to control Claude's response format."""
    print("\n=== Prefill Pattern ===")

    # Prefill forces Claude to start response in specific way
    message = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=512,
        messages=[
            {"role": "user", "content": "List 3 benefits of exercise"},
            {
                "role": "assistant",
                "content": "Here are the benefits:\n1."  # Prefill
            }
        ]
    )

    print("Prefilled response:")
    print("Here are the benefits:\n1." + message.content[0].text)


# ============================================================================
# Main Execution
# ============================================================================

def main():
    """Run all examples."""
    examples = [
        ("Basic Message", basic_message_example),
        ("XML-Structured Prompt", xml_structured_prompt),
        ("Tool Use", tool_use_example),
        ("Multi-Turn Conversation", multi_turn_conversation),
        ("Streaming", streaming_example),
        ("Prompt Caching", prompt_caching_example),
        ("Extended Thinking", extended_thinking_example),
        ("Structured Extraction", structured_extraction_example),
        ("RAG with Citations", rag_with_citations),
        ("Prefill Pattern", prefill_pattern),
    ]

    print("Anthropic Claude Examples")
    print("=" * 60)

    for name, example_func in examples:
        try:
            example_func()
        except Exception as e:
            print(f"\nError in {name}: {e}")

    print("\n" + "=" * 60)
    print("All examples completed!")


if __name__ == "__main__":
    # Check for API key
    if not os.environ.get("ANTHROPIC_API_KEY"):
        print("Error: ANTHROPIC_API_KEY environment variable not set")
        print("Set it with: export ANTHROPIC_API_KEY='your-api-key'")
        exit(1)

    main()

"""
LangChain Prompt Engineering Examples

Demonstrates LangChain patterns:
- PromptTemplates
- ChatPromptTemplates
- Few-shot examples
- Output parsers
- Chain composition (LCEL)

Installation:
    pip install langchain langchain-openai langchain-anthropic

Usage:
    export OPENAI_API_KEY="your-api-key"
    export ANTHROPIC_API_KEY="your-api-key"  # Optional
    python langchain-examples.py
"""

import os
from typing import List, Dict
from langchain_core.prompts import (
    PromptTemplate,
    ChatPromptTemplate,
    FewShotPromptTemplate,
    MessagesPlaceholder,
)
from langchain_core.output_parsers import (
    StrOutputParser,
    JsonOutputParser,
    PydanticOutputParser,
)
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain_openai import ChatOpenAI


# Initialize LLM
llm = ChatOpenAI(
    model="gpt-4",
    temperature=0.7,
    api_key=os.environ.get("OPENAI_API_KEY")
)


# ============================================================================
# Example 1: Basic PromptTemplate
# ============================================================================

def basic_prompt_template():
    """Simple string-based prompt template."""
    print("\n=== Basic PromptTemplate ===")

    template = PromptTemplate.from_template(
        "Translate the following text to {language}: {text}"
    )

    prompt = template.format(language="French", text="Hello, how are you?")
    print(f"Formatted prompt: {prompt}")

    # Use with LLM
    response = llm.invoke(prompt)
    print(f"Response: {response.content}")


# ============================================================================
# Example 2: ChatPromptTemplate
# ============================================================================

def chat_prompt_template():
    """Chat-based prompt template with system and user messages."""
    print("\n=== ChatPromptTemplate ===")

    template = ChatPromptTemplate.from_messages([
        ("system", "You are a {role} who {style}."),
        ("user", "{task}")
    ])

    messages = template.format_messages(
        role="creative writer",
        style="uses vivid imagery and metaphors",
        task="Write a short description of a sunset."
    )

    print("Formatted messages:")
    for msg in messages:
        print(f"  {msg.type}: {msg.content}")

    response = llm.invoke(messages)
    print(f"\nResponse: {response.content}")


# ============================================================================
# Example 3: Few-Shot Prompting
# ============================================================================

def few_shot_example():
    """Few-shot learning with examples."""
    print("\n=== Few-Shot Prompting ===")

    # Define examples
    examples = [
        {
            "input": "happy",
            "output": "😊"
        },
        {
            "input": "sad",
            "output": "😢"
        },
        {
            "input": "excited",
            "output": "🎉"
        }
    ]

    # Example formatter
    example_template = PromptTemplate(
        input_variables=["input", "output"],
        template="Emotion: {input}\nEmoji: {output}"
    )

    # Prefix and suffix
    prefix = "Convert emotions to emojis:\n\n"
    suffix = "\nEmotion: {input}\nEmoji:"

    # Create few-shot prompt
    few_shot_prompt = FewShotPromptTemplate(
        examples=examples,
        example_prompt=example_template,
        prefix=prefix,
        suffix=suffix,
        input_variables=["input"]
    )

    # Format and invoke
    prompt = few_shot_prompt.format(input="surprised")
    print(f"Formatted prompt:\n{prompt}\n")

    response = llm.invoke(prompt)
    print(f"Response: {response.content}")


# ============================================================================
# Example 4: Output Parsers - String Parser
# ============================================================================

def string_output_parser():
    """Parse string output."""
    print("\n=== String Output Parser ===")

    template = ChatPromptTemplate.from_messages([
        ("system", "You are a helpful assistant."),
        ("user", "{question}")
    ])

    parser = StrOutputParser()

    # Create chain using LCEL
    chain = template | llm | parser

    result = chain.invoke({"question": "What is the capital of France?"})
    print(f"Parsed result (string): {result}")
    print(f"Type: {type(result)}")


# ============================================================================
# Example 5: Output Parsers - JSON Parser
# ============================================================================

def json_output_parser():
    """Parse JSON output."""
    print("\n=== JSON Output Parser ===")

    parser = JsonOutputParser()

    template = ChatPromptTemplate.from_messages([
        ("system", "Extract information as JSON."),
        ("user", "{query}\n\n{format_instructions}")
    ])

    chain = template | llm | parser

    result = chain.invoke({
        "query": "Extract name and age: 'John is 25 years old'",
        "format_instructions": parser.get_format_instructions()
    })

    print(f"Parsed result (dict): {result}")
    print(f"Type: {type(result)}")


# ============================================================================
# Example 6: Output Parsers - Pydantic Parser
# ============================================================================

def pydantic_output_parser():
    """Parse output into Pydantic model."""
    print("\n=== Pydantic Output Parser ===")

    try:
        from pydantic import BaseModel, Field

        class Person(BaseModel):
            name: str = Field(description="Person's name")
            age: int = Field(description="Person's age")
            occupation: str = Field(description="Person's occupation")

        parser = PydanticOutputParser(pydantic_object=Person)

        template = ChatPromptTemplate.from_messages([
            ("system", "Extract person information."),
            ("user", "{query}\n\n{format_instructions}")
        ])

        chain = template | llm | parser

        result = chain.invoke({
            "query": "Parse: 'Alice is a 30-year-old engineer'",
            "format_instructions": parser.get_format_instructions()
        })

        print(f"Parsed result: {result}")
        print(f"Name: {result.name}")
        print(f"Age: {result.age}")
        print(f"Occupation: {result.occupation}")

    except ImportError:
        print("Install pydantic: pip install pydantic")


# ============================================================================
# Example 7: LCEL Chain Composition
# ============================================================================

def lcel_chain_composition():
    """Compose chains using LangChain Expression Language."""
    print("\n=== LCEL Chain Composition ===")

    # Step 1: Generate a topic
    topic_template = ChatPromptTemplate.from_template(
        "Suggest a specific topic related to: {subject}"
    )

    # Step 2: Write about the topic
    writing_template = ChatPromptTemplate.from_template(
        "Write a 2-sentence paragraph about: {topic}"
    )

    # Step 3: Summarize
    summary_template = ChatPromptTemplate.from_template(
        "Summarize this in one sentence: {paragraph}"
    )

    # Compose chain
    chain = (
        topic_template
        | llm
        | StrOutputParser()
        | (lambda topic: {"topic": topic})
        | writing_template
        | llm
        | StrOutputParser()
        | (lambda paragraph: {"paragraph": paragraph})
        | summary_template
        | llm
        | StrOutputParser()
    )

    result = chain.invoke({"subject": "artificial intelligence"})
    print(f"Final summary: {result}")


# ============================================================================
# Example 8: Conversation Memory
# ============================================================================

def conversation_with_memory():
    """Maintain conversation history."""
    print("\n=== Conversation with Memory ===")

    template = ChatPromptTemplate.from_messages([
        ("system", "You are a helpful assistant."),
        MessagesPlaceholder(variable_name="history"),
        ("user", "{input}")
    ])

    chain = template | llm | StrOutputParser()

    # Conversation history
    history = []

    def chat(user_input: str) -> str:
        """Send message and update history."""
        response = chain.invoke({
            "history": history,
            "input": user_input
        })

        # Update history
        history.append(HumanMessage(content=user_input))
        history.append(AIMessage(content=response))

        return response

    # Turn 1
    response1 = chat("My name is Alice and I love Python.")
    print(f"User: My name is Alice and I love Python.")
    print(f"Assistant: {response1}\n")

    # Turn 2
    response2 = chat("What's my name and what language do I like?")
    print(f"User: What's my name and what language do I like?")
    print(f"Assistant: {response2}")


# ============================================================================
# Example 9: Conditional Chain (Routing)
# ============================================================================

def conditional_routing():
    """Route to different chains based on input."""
    print("\n=== Conditional Routing ===")

    from langchain_core.runnables import RunnableBranch

    # Classification chain
    classify_template = ChatPromptTemplate.from_template(
        """Classify this query into ONE category:
- technical
- billing
- general

Query: {query}

Category (one word):"""
    )

    classify_chain = classify_template | llm | StrOutputParser()

    # Technical support chain
    technical_template = ChatPromptTemplate.from_template(
        "Provide technical support for: {query}"
    )
    technical_chain = technical_template | llm | StrOutputParser()

    # Billing chain
    billing_template = ChatPromptTemplate.from_template(
        "Provide billing assistance for: {query}"
    )
    billing_chain = billing_template | llm | StrOutputParser()

    # General chain
    general_template = ChatPromptTemplate.from_template(
        "Provide general assistance for: {query}"
    )
    general_chain = general_template | llm | StrOutputParser()

    # Route based on classification
    def route_query(query: str):
        category = classify_chain.invoke({"query": query}).strip().lower()
        print(f"Classified as: {category}")

        if "technical" in category:
            return technical_chain.invoke({"query": query})
        elif "billing" in category:
            return billing_chain.invoke({"query": query})
        else:
            return general_chain.invoke({"query": query})

    # Test
    query = "How do I reset my password?"
    response = route_query(query)
    print(f"\nQuery: {query}")
    print(f"Response: {response}")


# ============================================================================
# Example 10: Parallel Chain Execution
# ============================================================================

def parallel_chains():
    """Execute multiple chains in parallel."""
    print("\n=== Parallel Chain Execution ===")

    from langchain_core.runnables import RunnableParallel

    # Define parallel tasks
    pros_template = ChatPromptTemplate.from_template(
        "List 3 pros of: {topic}"
    )
    cons_template = ChatPromptTemplate.from_template(
        "List 3 cons of: {topic}"
    )
    summary_template = ChatPromptTemplate.from_template(
        "Summarize in one sentence: {topic}"
    )

    # Create parallel chain
    parallel_chain = RunnableParallel(
        pros=pros_template | llm | StrOutputParser(),
        cons=cons_template | llm | StrOutputParser(),
        summary=summary_template | llm | StrOutputParser()
    )

    result = parallel_chain.invoke({"topic": "remote work"})

    print("Parallel results:")
    print(f"\nPros:\n{result['pros']}")
    print(f"\nCons:\n{result['cons']}")
    print(f"\nSummary:\n{result['summary']}")


# ============================================================================
# Example 11: Custom Prompt Template
# ============================================================================

def custom_prompt_template():
    """Create custom prompt template with validation."""
    print("\n=== Custom Prompt Template ===")

    class EmailPromptTemplate(PromptTemplate):
        """Custom template for email generation."""

        def format(self, **kwargs) -> str:
            # Validate inputs
            if "recipient" not in kwargs:
                raise ValueError("recipient is required")
            if "tone" not in kwargs:
                kwargs["tone"] = "professional"

            # Add default signature
            if "signature" not in kwargs:
                kwargs["signature"] = "Best regards,\n[Your Name]"

            return super().format(**kwargs)

    template = EmailPromptTemplate(
        input_variables=["recipient", "subject", "tone", "signature"],
        template="""
Write a {tone} email to {recipient} with the subject: {subject}

Email:

{signature}
"""
    )

    prompt = template.format(
        recipient="Sarah",
        subject="Project Update",
        tone="friendly"
    )

    print(f"Formatted prompt:\n{prompt}")

    response = llm.invoke(prompt)
    print(f"\nGenerated email:\n{response.content}")


# ============================================================================
# Example 12: Multi-Step Reasoning Chain
# ============================================================================

def multi_step_reasoning():
    """Chain for complex multi-step reasoning."""
    print("\n=== Multi-Step Reasoning Chain ===")

    # Step 1: Break down the problem
    breakdown_template = ChatPromptTemplate.from_template(
        """Break this problem into 3-5 logical steps:

Problem: {problem}

Steps:"""
    )

    # Step 2: Solve each step
    solve_template = ChatPromptTemplate.from_template(
        """Original problem: {problem}

Steps to solve:
{steps}

Solve each step and provide the final answer:"""
    )

    # Create chain
    chain = (
        breakdown_template
        | llm
        | StrOutputParser()
        | (lambda steps: {"problem": "PLACEHOLDER", "steps": steps})
        | solve_template
        | llm
        | StrOutputParser()
    )

    # We need to manually handle the problem variable
    problem = "If a train travels 120 miles in 2 hours, then speeds up to travel 180 miles in 2 hours, what is its average speed for the entire journey?"

    # Step 1
    steps = (breakdown_template | llm | StrOutputParser()).invoke({"problem": problem})
    print(f"Steps:\n{steps}\n")

    # Step 2
    solution = (solve_template | llm | StrOutputParser()).invoke({
        "problem": problem,
        "steps": steps
    })
    print(f"Solution:\n{solution}")


# ============================================================================
# Main Execution
# ============================================================================

def main():
    """Run all examples."""
    examples = [
        ("Basic PromptTemplate", basic_prompt_template),
        ("ChatPromptTemplate", chat_prompt_template),
        ("Few-Shot Example", few_shot_example),
        ("String Output Parser", string_output_parser),
        ("JSON Output Parser", json_output_parser),
        ("Pydantic Output Parser", pydantic_output_parser),
        ("LCEL Chain Composition", lcel_chain_composition),
        ("Conversation with Memory", conversation_with_memory),
        ("Conditional Routing", conditional_routing),
        ("Parallel Chains", parallel_chains),
        ("Custom Prompt Template", custom_prompt_template),
        ("Multi-Step Reasoning", multi_step_reasoning),
    ]

    print("LangChain Prompt Engineering Examples")
    print("=" * 60)

    for name, example_func in examples:
        try:
            example_func()
        except Exception as e:
            print(f"\nError in {name}: {e}")

    print("\n" + "=" * 60)
    print("All examples completed!")


if __name__ == "__main__":
    # Check for API key
    if not os.environ.get("OPENAI_API_KEY"):
        print("Error: OPENAI_API_KEY environment variable not set")
        print("Set it with: export OPENAI_API_KEY='your-api-key'")
        exit(1)

    main()

"""
OpenAI Prompt Engineering Examples

Demonstrates OpenAI-specific features:
- Chat completions API
- Function calling
- JSON mode
- Structured outputs
- Vision prompts (GPT-4o)
- Streaming

Installation:
    pip install openai

Usage:
    export OPENAI_API_KEY="your-api-key"
    python openai-examples.py
"""

import os
import json
from openai import OpenAI
from typing import List, Dict, Any


# Initialize client
client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))


# ============================================================================
# Example 1: Basic Chat Completion
# ============================================================================

def basic_completion():
    """Simple chat completion with system and user messages."""
    print("\n=== Basic Chat Completion ===")

    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {
                "role": "system",
                "content": "You are a helpful assistant specialized in Python programming."
            },
            {
                "role": "user",
                "content": "Write a function to check if a string is a palindrome."
            }
        ],
        temperature=0.7,
        max_tokens=500
    )

    print(response.choices[0].message.content)
    print(f"\nTokens used: {response.usage.total_tokens}")


# ============================================================================
# Example 2: JSON Mode (Reliable JSON Output)
# ============================================================================

def json_mode_example():
    """Use JSON mode for guaranteed valid JSON output."""
    print("\n=== JSON Mode Example ===")

    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {
                "role": "system",
                "content": "You are a data extraction assistant. Extract information as JSON."
            },
            {
                "role": "user",
                "content": """
Extract the following information as JSON:

"Sarah Johnson, a 28-year-old software engineer from San Francisco,
enjoys hiking and photography. Contact: sarah@email.com"

Include: name, age, occupation, city, hobbies (array), email
"""
            }
        ],
        response_format={"type": "json_object"},
        temperature=0
    )

    json_output = json.loads(response.choices[0].message.content)
    print("Extracted JSON:")
    print(json.dumps(json_output, indent=2))


# ============================================================================
# Example 3: Function Calling
# ============================================================================

def function_calling_example():
    """Demonstrate function calling for tool use."""
    print("\n=== Function Calling Example ===")

    # Define available functions
    tools = [
        {
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "Get the current weather for a specific location",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "City and state, e.g., 'San Francisco, CA'"
                        },
                        "units": {
                            "type": "string",
                            "enum": ["celsius", "fahrenheit"],
                            "description": "Temperature unit"
                        }
                    },
                    "required": ["location"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "calculate",
                "description": "Perform a mathematical calculation",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "expression": {
                            "type": "string",
                            "description": "Mathematical expression, e.g., '2 + 2 * 3'"
                        }
                    },
                    "required": ["expression"]
                }
            }
        }
    ]

    # Initial request
    messages = [
        {"role": "user", "content": "What's the weather in Tokyo and what's 15 * 23?"}
    ]

    response = client.chat.completions.create(
        model="gpt-4",
        messages=messages,
        tools=tools,
        tool_choice="auto"
    )

    response_message = response.choices[0].message
    print(f"Assistant wants to call {len(response_message.tool_calls or [])} tool(s)")

    # Process tool calls
    if response_message.tool_calls:
        for tool_call in response_message.tool_calls:
            function_name = tool_call.function.name
            function_args = json.loads(tool_call.function.arguments)

            print(f"\nTool: {function_name}")
            print(f"Arguments: {function_args}")

            # Simulate function execution
            if function_name == "get_weather":
                result = json.dumps({
                    "temperature": 18,
                    "condition": "Partly cloudy",
                    "humidity": 65
                })
            elif function_name == "calculate":
                result = json.dumps({
                    "result": eval(function_args["expression"])
                })

            print(f"Result: {result}")


# ============================================================================
# Example 4: Multi-Turn Conversation with Memory
# ============================================================================

def conversation_example():
    """Demonstrate conversation history management."""
    print("\n=== Multi-Turn Conversation ===")

    conversation = []

    def chat(user_message: str) -> str:
        """Send message and get response."""
        conversation.append({"role": "user", "content": user_message})

        response = client.chat.completions.create(
            model="gpt-4",
            messages=conversation,
            temperature=0.7,
            max_tokens=200
        )

        assistant_message = response.choices[0].message.content
        conversation.append({"role": "assistant", "content": assistant_message})

        return assistant_message

    # Turn 1
    user_msg_1 = "My favorite color is blue and I work as a teacher."
    response_1 = chat(user_msg_1)
    print(f"User: {user_msg_1}")
    print(f"Assistant: {response_1}\n")

    # Turn 2
    user_msg_2 = "What's my favorite color and what do I do for work?"
    response_2 = chat(user_msg_2)
    print(f"User: {user_msg_2}")
    print(f"Assistant: {response_2}")


# ============================================================================
# Example 5: Streaming Responses
# ============================================================================

def streaming_example():
    """Stream response token-by-token."""
    print("\n=== Streaming Example ===")
    print("Assistant (streaming): ", end="", flush=True)

    stream = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "user", "content": "Write a haiku about programming"}
        ],
        stream=True,
        temperature=0.8
    )

    for chunk in stream:
        if chunk.choices[0].delta.content is not None:
            print(chunk.choices[0].delta.content, end="", flush=True)

    print("\n")


# ============================================================================
# Example 6: Few-Shot Learning
# ============================================================================

def few_shot_example():
    """Use few-shot examples to guide model behavior."""
    print("\n=== Few-Shot Learning ===")

    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {
                "role": "system",
                "content": "Classify the sentiment of movie reviews as positive, negative, or neutral."
            },
            # Example 1
            {
                "role": "user",
                "content": "This movie was absolutely fantastic! Best film of the year."
            },
            {
                "role": "assistant",
                "content": "Sentiment: positive"
            },
            # Example 2
            {
                "role": "user",
                "content": "Terrible acting and boring plot. Waste of time."
            },
            {
                "role": "assistant",
                "content": "Sentiment: negative"
            },
            # Example 3
            {
                "role": "user",
                "content": "It was okay. Nothing special but not terrible either."
            },
            {
                "role": "assistant",
                "content": "Sentiment: neutral"
            },
            # Actual query
            {
                "role": "user",
                "content": "The cinematography was beautiful but the story dragged on too long."
            }
        ],
        temperature=0,
        max_tokens=50
    )

    print(response.choices[0].message.content)


# ============================================================================
# Example 7: Chain-of-Thought Prompting
# ============================================================================

def chain_of_thought_example():
    """Use chain-of-thought for complex reasoning."""
    print("\n=== Chain-of-Thought Example ===")

    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {
                "role": "system",
                "content": "You are a math tutor. Always show your step-by-step reasoning."
            },
            {
                "role": "user",
                "content": """
A store sells apples for $0.50 each and oranges for $0.75 each.
If John buys 12 apples and 8 oranges, and he pays with a $20 bill,
how much change does he receive?

Think step-by-step:
"""
            }
        ],
        temperature=0,
        max_tokens=500
    )

    print(response.choices[0].message.content)


# ============================================================================
# Example 8: Temperature Variations
# ============================================================================

def temperature_comparison():
    """Compare outputs at different temperature settings."""
    print("\n=== Temperature Comparison ===")

    prompt = "Describe a sunset in one sentence."
    temperatures = [0, 0.5, 1.0, 1.5]

    for temp in temperatures:
        response = client.chat.completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": prompt}],
            temperature=temp,
            max_tokens=100
        )

        print(f"\nTemperature {temp}:")
        print(response.choices[0].message.content)


# ============================================================================
# Example 9: System Prompt Engineering
# ============================================================================

def system_prompt_comparison():
    """Compare different system prompts for the same task."""
    print("\n=== System Prompt Comparison ===")

    user_question = "Should I invest in cryptocurrency?"

    # System prompt 1: Neutral advisor
    print("\n1. Neutral Financial Advisor:")
    response1 = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {
                "role": "system",
                "content": """You are a neutral financial advisor.
Provide balanced information about investment risks and opportunities.
Always mention that this is not personalized financial advice."""
            },
            {"role": "user", "content": user_question}
        ],
        temperature=0.5,
        max_tokens=200
    )
    print(response1.choices[0].message.content)

    # System prompt 2: Conservative advisor
    print("\n2. Conservative Financial Advisor:")
    response2 = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {
                "role": "system",
                "content": """You are a conservative financial advisor who prioritizes
capital preservation and risk mitigation. You are skeptical of high-risk investments."""
            },
            {"role": "user", "content": user_question}
        ],
        temperature=0.5,
        max_tokens=200
    )
    print(response2.choices[0].message.content)


# ============================================================================
# Example 10: Structured Outputs (Pydantic Models)
# ============================================================================

def structured_outputs_example():
    """Use OpenAI's structured outputs with Pydantic."""
    print("\n=== Structured Outputs with Pydantic ===")

    try:
        from pydantic import BaseModel
        from typing import List

        class UserProfile(BaseModel):
            name: str
            age: int
            email: str
            interests: List[str]
            occupation: str

        response = client.beta.chat.completions.parse(
            model="gpt-4o-2024-08-06",  # Requires model with structured outputs
            messages=[
                {
                    "role": "system",
                    "content": "Extract user profile information in the specified format."
                },
                {
                    "role": "user",
                    "content": """
Parse this bio:
"I'm Alex Chen, 32 years old, working as a data scientist.
I love hiking, photography, and reading sci-fi novels.
You can email me at alex.chen@example.com"
"""
                }
            ],
            response_format=UserProfile
        )

        user_profile = response.choices[0].message.parsed
        print("Parsed Profile:")
        print(f"Name: {user_profile.name}")
        print(f"Age: {user_profile.age}")
        print(f"Email: {user_profile.email}")
        print(f"Occupation: {user_profile.occupation}")
        print(f"Interests: {', '.join(user_profile.interests)}")

    except ImportError:
        print("Note: Install pydantic for structured outputs: pip install pydantic")
    except Exception as e:
        print(f"Structured outputs not available: {e}")
        print("Falling back to JSON mode...")

        # Fallback to regular JSON mode
        response = client.chat.completions.create(
            model="gpt-4",
            messages=[
                {
                    "role": "system",
                    "content": "Extract user profile as JSON with fields: name, age, email, interests (array), occupation"
                },
                {
                    "role": "user",
                    "content": """
Parse this bio:
"I'm Alex Chen, 32 years old, working as a data scientist.
I love hiking, photography, and reading sci-fi novels.
You can email me at alex.chen@example.com"
"""
                }
            ],
            response_format={"type": "json_object"}
        )

        profile = json.loads(response.choices[0].message.content)
        print(json.dumps(profile, indent=2))


# ============================================================================
# Example 11: Vision (GPT-4o Image Analysis)
# ============================================================================

def vision_example():
    """Analyze images with GPT-4o (vision capabilities)."""
    print("\n=== Vision Example (GPT-4o) ===")

    # Using a publicly accessible image URL
    image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/320px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "Describe this image in detail. What's the setting, mood, and key elements?"
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": image_url
                        }
                    }
                ]
            }
        ],
        max_tokens=300
    )

    print(response.choices[0].message.content)


# ============================================================================
# Example 12: Token Usage Tracking
# ============================================================================

def token_usage_example():
    """Track token usage and estimate costs."""
    print("\n=== Token Usage Tracking ===")

    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {
                "role": "user",
                "content": "Explain quantum computing in 2 paragraphs."
            }
        ]
    )

    usage = response.usage
    print(f"Prompt tokens: {usage.prompt_tokens}")
    print(f"Completion tokens: {usage.completion_tokens}")
    print(f"Total tokens: {usage.total_tokens}")

    # Cost estimation (approximate prices as of 2024)
    input_cost_per_1k = 0.01  # $0.01 per 1K input tokens for GPT-4
    output_cost_per_1k = 0.03  # $0.03 per 1K output tokens for GPT-4

    estimated_cost = (
        (usage.prompt_tokens / 1000) * input_cost_per_1k +
        (usage.completion_tokens / 1000) * output_cost_per_1k
    )

    print(f"Estimated cost: ${estimated_cost:.6f}")


# ============================================================================
# Main Execution
# ============================================================================

def main():
    """Run all examples."""
    examples = [
        ("Basic Completion", basic_completion),
        ("JSON Mode", json_mode_example),
        ("Function Calling", function_calling_example),
        ("Multi-Turn Conversation", conversation_example),
        ("Streaming", streaming_example),
        ("Few-Shot Learning", few_shot_example),
        ("Chain-of-Thought", chain_of_thought_example),
        ("Temperature Comparison", temperature_comparison),
        ("System Prompt Comparison", system_prompt_comparison),
        ("Structured Outputs", structured_outputs_example),
        ("Vision (GPT-4o)", vision_example),
        ("Token Usage Tracking", token_usage_example),
    ]

    print("OpenAI Prompt Engineering Examples")
    print("=" * 60)

    for name, example_func in examples:
        try:
            example_func()
        except Exception as e:
            print(f"\nError in {name}: {e}")

    print("\n" + "=" * 60)
    print("All examples completed!")


if __name__ == "__main__":
    # Check for API key
    if not os.environ.get("OPENAI_API_KEY"):
        print("Error: OPENAI_API_KEY environment variable not set")
        print("Set it with: export OPENAI_API_KEY='your-api-key'")
        exit(1)

    main()

"""
Complete RAG (Retrieval-Augmented Generation) Implementation

A production-ready RAG system demonstrating:
- Document chunking
- Embedding generation
- Vector search
- Context assembly
- Answer generation with citations
- Full working example

Installation:
    pip install openai chromadb pypdf sentence-transformers

Usage:
    export OPENAI_API_KEY="your-api-key"
    python rag-complete-example.py
"""

import os
import hashlib
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass
from openai import OpenAI


# ============================================================================
# Configuration
# ============================================================================

@dataclass
class RAGConfig:
    """RAG system configuration."""
    chunk_size: int = 512
    chunk_overlap: int = 100
    top_k: int = 5
    model: str = "gpt-4"
    embedding_model: str = "text-embedding-3-small"
    temperature: float = 0.3


# ============================================================================
# Document Processing
# ============================================================================

class DocumentChunker:
    """Split documents into chunks for embedding."""

    def __init__(self, chunk_size: int = 512, overlap: int = 100):
        self.chunk_size = chunk_size
        self.overlap = overlap

    def chunk_text(self, text: str, metadata: Optional[Dict] = None) -> List[Dict]:
        """
        Split text into overlapping chunks.

        Args:
            text: Text to chunk
            metadata: Optional metadata (source, date, etc.)

        Returns:
            List of chunk dictionaries with text and metadata
        """
        chunks = []
        start = 0

        while start < len(text):
            end = min(start + self.chunk_size, len(text))

            # Try to break at sentence boundary
            if end < len(text):
                # Look for sentence ending
                for punctuation in ['. ', '! ', '? ', '\n\n']:
                    last_punct = text.rfind(punctuation, start, end)
                    if last_punct > start:
                        end = last_punct + len(punctuation)
                        break

            chunk_text = text[start:end].strip()

            if chunk_text:
                chunk = {
                    'text': chunk_text,
                    'start_pos': start,
                    'end_pos': end,
                    'chunk_id': len(chunks)
                }

                # Add metadata if provided
                if metadata:
                    chunk.update(metadata)

                chunks.append(chunk)

            start = end - self.overlap

        return chunks


# ============================================================================
# Vector Store (using ChromaDB)
# ============================================================================

class VectorStore:
    """Vector database for storing and retrieving embeddings."""

    def __init__(self, embedding_model: str = "text-embedding-3-small"):
        self.client = OpenAI()
        self.embedding_model = embedding_model
        self.documents: List[Dict] = []
        self.embeddings: List[List[float]] = []

    def add_documents(self, chunks: List[Dict]):
        """
        Add documents to the vector store.

        Args:
            chunks: List of document chunks with text and metadata
        """
        print(f"Adding {len(chunks)} chunks to vector store...")

        for chunk in chunks:
            # Generate embedding
            embedding = self._get_embedding(chunk['text'])

            # Store
            self.documents.append(chunk)
            self.embeddings.append(embedding)

        print(f"Vector store now contains {len(self.documents)} chunks")

    def _get_embedding(self, text: str) -> List[float]:
        """Generate embedding for text."""
        response = self.client.embeddings.create(
            model=self.embedding_model,
            input=text
        )
        return response.data[0].embedding

    def search(self, query: str, top_k: int = 5) -> List[Tuple[Dict, float]]:
        """
        Search for most similar documents.

        Args:
            query: Search query
            top_k: Number of results to return

        Returns:
            List of (document, similarity_score) tuples
        """
        # Get query embedding
        query_embedding = self._get_embedding(query)

        # Calculate cosine similarity
        similarities = []
        for i, doc_embedding in enumerate(self.embeddings):
            similarity = self._cosine_similarity(query_embedding, doc_embedding)
            similarities.append((self.documents[i], similarity))

        # Sort by similarity and return top-k
        similarities.sort(key=lambda x: x[1], reverse=True)
        return similarities[:top_k]

    @staticmethod
    def _cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
        """Calculate cosine similarity between two vectors."""
        dot_product = sum(a * b for a, b in zip(vec1, vec2))
        magnitude1 = sum(a * a for a in vec1) ** 0.5
        magnitude2 = sum(b * b for b in vec2) ** 0.5
        return dot_product / (magnitude1 * magnitude2)


# ============================================================================
# RAG System
# ============================================================================

class RAGSystem:
    """Complete RAG system for question-answering."""

    def __init__(self, config: Optional[RAGConfig] = None):
        self.config = config or RAGConfig()
        self.client = OpenAI()
        self.chunker = DocumentChunker(
            chunk_size=self.config.chunk_size,
            overlap=self.config.chunk_overlap
        )
        self.vector_store = VectorStore(
            embedding_model=self.config.embedding_model
        )

    def add_documents(self, documents: List[Dict[str, str]]):
        """
        Add documents to the knowledge base.

        Args:
            documents: List of dicts with 'text' and optional metadata
        """
        all_chunks = []

        for doc in documents:
            # Extract text and metadata
            text = doc['text']
            metadata = {k: v for k, v in doc.items() if k != 'text'}

            # Chunk document
            chunks = self.chunker.chunk_text(text, metadata)
            all_chunks.extend(chunks)

        # Add to vector store
        self.vector_store.add_documents(all_chunks)

    def query(self, question: str, return_sources: bool = True) -> Dict:
        """
        Answer a question using RAG.

        Args:
            question: User question
            return_sources: Whether to return source documents

        Returns:
            Dictionary with answer and optional sources
        """
        # Step 1: Retrieve relevant documents
        retrieved_docs = self.vector_store.search(
            question,
            top_k=self.config.top_k
        )

        # Step 2: Format context
        context = self._format_context(retrieved_docs)

        # Step 3: Generate answer
        answer = self._generate_answer(question, context)

        result = {'answer': answer}

        if return_sources:
            result['sources'] = [
                {
                    'text': doc['text'][:200] + '...',  # Preview
                    'source': doc.get('source', 'unknown'),
                    'similarity': score
                }
                for doc, score in retrieved_docs
            ]

        return result

    def _format_context(self, retrieved_docs: List[Tuple[Dict, float]]) -> str:
        """Format retrieved documents as context."""
        context_parts = []

        for i, (doc, score) in enumerate(retrieved_docs):
            source = doc.get('source', 'Unknown')
            text = doc['text']

            context_parts.append(f"""
[Document {i+1}]
Source: {source}
Relevance: {score:.2f}
Content: {text}
""")

        return "\n".join(context_parts)

    def _generate_answer(self, question: str, context: str) -> str:
        """Generate answer using LLM."""
        system_prompt = """
You are a helpful AI assistant that answers questions based on provided context.

IMPORTANT RULES:
1. Answer ONLY using information from the provided documents
2. Cite sources using [Document N] notation
3. If the answer is not in the documents, say "I don't have enough information to answer this question."
4. Be concise but complete
5. Use direct quotes when appropriate
"""

        user_prompt = f"""
Context:
{context}

Question: {question}

Answer (with citations):
"""

        response = self.client.chat.completions.create(
            model=self.config.model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            temperature=self.config.temperature,
            max_tokens=1024
        )

        return response.choices[0].message.content


# ============================================================================
# Advanced RAG Features
# ============================================================================

class AdvancedRAG(RAGSystem):
    """RAG system with advanced features."""

    def query_with_rerank(self, question: str) -> Dict:
        """
        Query with re-ranking for better relevance.

        Uses a two-stage retrieval:
        1. Retrieve top-K candidates (e.g., 20)
        2. Re-rank to select best top-N (e.g., 5)
        """
        # Stage 1: Retrieve more candidates
        candidates = self.vector_store.search(question, top_k=20)

        # Stage 2: Re-rank (simplified - in production use cross-encoder)
        reranked = self._rerank(question, candidates)

        # Use top-k after reranking
        top_docs = reranked[:self.config.top_k]

        # Generate answer
        context = self._format_context(top_docs)
        answer = self._generate_answer(question, context)

        return {
            'answer': answer,
            'sources': [
                {
                    'text': doc['text'][:200] + '...',
                    'source': doc.get('source', 'unknown'),
                    'similarity': score
                }
                for doc, score in top_docs
            ]
        }

    def _rerank(
        self,
        query: str,
        candidates: List[Tuple[Dict, float]]
    ) -> List[Tuple[Dict, float]]:
        """
        Re-rank candidates for better relevance.

        Simplified reranking - in production, use a cross-encoder model.
        """
        # For this example, we'll just boost scores based on exact keyword matches
        query_keywords = set(query.lower().split())

        reranked = []
        for doc, score in candidates:
            # Calculate keyword overlap
            doc_keywords = set(doc['text'].lower().split())
            overlap = len(query_keywords & doc_keywords)

            # Boost score
            boosted_score = score * (1 + 0.1 * overlap)
            reranked.append((doc, boosted_score))

        # Sort by boosted score
        reranked.sort(key=lambda x: x[1], reverse=True)
        return reranked

    def conversational_query(
        self,
        question: str,
        conversation_history: List[Dict]
    ) -> Dict:
        """
        Query with conversation history for context.

        Args:
            question: Current question
            conversation_history: List of previous Q&A pairs

        Returns:
            Answer with updated conversation history
        """
        # Retrieve relevant documents
        retrieved_docs = self.vector_store.search(question, top_k=self.config.top_k)
        context = self._format_context(retrieved_docs)

        # Build conversation with history
        messages = [
            {
                "role": "system",
                "content": """You are a helpful AI assistant that answers questions
based on provided context. Maintain context from previous conversation."""
            }
        ]

        # Add conversation history
        for item in conversation_history[-3:]:  # Last 3 turns
            messages.append({"role": "user", "content": item['question']})
            messages.append({"role": "assistant", "content": item['answer']})

        # Add current query
        messages.append({
            "role": "user",
            "content": f"""
Context:
{context}

Question: {question}

Answer (with citations):
"""
        })

        response = self.client.chat.completions.create(
            model=self.config.model,
            messages=messages,
            temperature=self.config.temperature
        )

        answer = response.choices[0].message.content

        return {
            'answer': answer,
            'sources': [
                {'text': doc['text'][:200] + '...', 'source': doc.get('source', 'unknown')}
                for doc, _ in retrieved_docs
            ]
        }


# ============================================================================
# Example Usage
# ============================================================================

def demo_basic_rag():
    """Demonstrate basic RAG functionality."""
    print("\n" + "=" * 60)
    print("DEMO: Basic RAG System")
    print("=" * 60)

    # Initialize RAG system
    rag = RAGSystem()

    # Add sample documents
    documents = [
        {
            'text': """
Python is a high-level, interpreted programming language known for its
clear syntax and readability. It was created by Guido van Rossum and
first released in 1991. Python supports multiple programming paradigms
including procedural, object-oriented, and functional programming.
""",
            'source': 'python_intro.txt',
            'date': '2024-01-15'
        },
        {
            'text': """
Machine learning is a subset of artificial intelligence that enables
systems to learn and improve from experience without being explicitly
programmed. It focuses on developing computer programs that can access
data and use it to learn for themselves.
""",
            'source': 'ml_basics.txt',
            'date': '2024-01-20'
        },
        {
            'text': """
Neural networks are computing systems inspired by biological neural
networks in animal brains. They consist of interconnected nodes (neurons)
organized in layers. Deep learning uses neural networks with many layers
to learn from large amounts of data.
""",
            'source': 'neural_networks.txt',
            'date': '2024-01-25'
        }
    ]

    print("\nAdding documents to knowledge base...")
    rag.add_documents(documents)

    # Query the system
    questions = [
        "Who created Python?",
        "What is machine learning?",
        "How do neural networks work?"
    ]

    for question in questions:
        print(f"\n{'─' * 60}")
        print(f"Question: {question}")
        print(f"{'─' * 60}")

        result = rag.query(question)

        print(f"\nAnswer:\n{result['answer']}")

        print("\nSources:")
        for i, source in enumerate(result['sources'], 1):
            print(f"\n{i}. {source['source']} (Similarity: {source['similarity']:.3f})")
            print(f"   {source['text']}")


def demo_advanced_rag():
    """Demonstrate advanced RAG features."""
    print("\n" + "=" * 60)
    print("DEMO: Advanced RAG with Re-ranking")
    print("=" * 60)

    # Initialize advanced RAG
    rag = AdvancedRAG()

    # Add technical documentation
    documents = [
        {
            'text': """
FastAPI is a modern, fast web framework for building APIs with Python 3.7+
based on standard Python type hints. It's one of the fastest Python frameworks
available, comparable to NodeJS and Go. Key features include automatic API
documentation, data validation using Pydantic, and async support.
""",
            'source': 'fastapi_docs.md'
        },
        {
            'text': """
Django is a high-level Python web framework that encourages rapid development
and clean, pragmatic design. It follows the model-template-view (MTV) architectural
pattern. Django includes an ORM, authentication system, and admin interface
out of the box.
""",
            'source': 'django_docs.md'
        },
        {
            'text': """
Flask is a lightweight WSGI web application framework in Python. It's designed
to make getting started quick and easy, with the ability to scale up to complex
applications. Flask is considered more Pythonic than Django because it's explicit
and doesn't make many decisions for you.
""",
            'source': 'flask_docs.md'
        }
    ]

    print("\nAdding documents...")
    rag.add_documents(documents)

    # Query with re-ranking
    question = "Which Python web framework is fastest?"
    print(f"\nQuestion: {question}")

    result = rag.query_with_rerank(question)

    print(f"\nAnswer:\n{result['answer']}")

    print("\nTop Sources (after re-ranking):")
    for i, source in enumerate(result['sources'], 1):
        print(f"{i}. {source['source']} (Score: {source['similarity']:.3f})")


def demo_conversational_rag():
    """Demonstrate conversational RAG."""
    print("\n" + "=" * 60)
    print("DEMO: Conversational RAG")
    print("=" * 60)

    rag = AdvancedRAG()

    # Add context
    documents = [
        {
            'text': """
React is a JavaScript library for building user interfaces, maintained by
Facebook. It uses a component-based architecture and a virtual DOM for
efficient updates. React introduced the concept of JSX, a syntax extension
that allows writing HTML-like code in JavaScript.
""",
            'source': 'react_guide.md'
        }
    ]

    rag.add_documents(documents)

    # Conversation
    conversation_history = []

    questions = [
        "What is React?",
        "Who maintains it?",
        "What is JSX?"
    ]

    for question in questions:
        print(f"\n{'─' * 60}")
        print(f"User: {question}")

        result = rag.conversational_query(question, conversation_history)

        print(f"Assistant: {result['answer']}")

        # Update history
        conversation_history.append({
            'question': question,
            'answer': result['answer']
        })


# ============================================================================
# Main
# ============================================================================

def main():
    """Run all demos."""
    print("Complete RAG System Demo")
    print("=" * 60)

    try:
        demo_basic_rag()
        demo_advanced_rag()
        demo_conversational_rag()

        print("\n" + "=" * 60)
        print("All demos completed successfully!")
        print("=" * 60)

    except Exception as e:
        print(f"\nError: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    # Check for API key
    if not os.environ.get("OPENAI_API_KEY"):
        print("Error: OPENAI_API_KEY environment variable not set")
        print("Set it with: export OPENAI_API_KEY='your-api-key'")
        exit(1)

    main()