11_type_hints.py

"""
Module 11: Type Hints and TypedDict
Learn: type annotations, Optional, TypedDict

In the advanced courses, you'll see:
    class AgentState(TypedDict):
        question: str
        sql_query: str
        error: str
    def run_bash(command: str) -> str: ...
"""

# --- Basic type hints ---
# They don't change behavior — they're documentation for humans and tools
name: str = "gemini-2.0-flash"
temperature: float = 0.7
max_tokens: int = 8000
is_active: bool = True

# Function with type hints
def truncate(text: str, max_len: int = 50) -> str:
    if len(text) <= max_len:
        return text
    return text[:max_len] + "..."

result = truncate("This is a very long text that needs to be shortened", 20)
print(f"Truncated: {result}")

# --- Collection type hints ---
scores: list[int] = [3, 5, 1, 4]
config: dict[str, str] = {"role": "user", "content": "hello"}
tools: list[dict[str, str]] = [
    {"name": "bash", "description": "Run shell commands"},
    {"name": "read_file", "description": "Read a file"},
]

def process_messages(messages: list[dict]) -> int:
    return len(messages)

print(f"Tool count: {len(tools)}")

# --- Optional and Union ---
from typing import Optional

def get_env(key: str, default: Optional[str] = None) -> Optional[str]:
    """Returns the value or None."""
    import os
    return os.getenv(key, default)

# Modern syntax (Python 3.10+) — you can use | instead of Optional
# These two are equivalent:
#   Optional[str]  ==  str | None
def get_env_modern(key: str, default: str | None = None) -> str | None:
    """Same function using modern union syntax."""
    import os
    return os.getenv(key, default)

# --- TypedDict (used in LangGraph state) ---
from typing import TypedDict

class AgentState(TypedDict):
    """This is exactly how the text-to-sql agent defines its state."""
    question: str
    sql_query: str
    query_result: str
    error: str
    iteration: int

# Create a state (it's just a dict with structure)
state: AgentState = {
    "question": "How many orders?",
    "sql_query": "",
    "query_result": "",
    "error": "",
    "iteration": 0,
}

# Use it like a normal dict
state["sql_query"] = "SELECT COUNT(*) FROM orders"
state["iteration"] += 1
print(f"\nAgent state: {state['question']}")
print(f"SQL: {state['sql_query']}")
print(f"Iteration: {state['iteration']}")

# 🎯 Exercise: Create a TypedDict for a ChatMessage with role, content, and timestamp