Lab — API Power Workshop

import os
from typing import Literal
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider

load_dotenv()

PROXY_URL = "https://litellm.6640.ucf.spencerlyon.com"

def get_model(model_name: str) -> OpenAIChatModel:
    """Create a model connection through our LiteLLM proxy."""
    return OpenAIChatModel(
        model_name,
        provider=OpenAIProvider(
            base_url=PROXY_URL,
            api_key=os.environ["CAP6640_API_KEY"],
        ),
    )

class MenuItem(BaseModel):
    """A single item on the menu."""
    name: str = Field(description="Name of the dish")
    description: str | None = Field(description="Description of the dish, if provided")
    price: float | None = Field(description="Price in dollars, or null if not listed")
    dietary_tags: list[Literal[
        "vegetarian", "vegan", "gluten-free", "spicy", "contains-nuts"
    ]] = Field(
        default_factory=list,
        description="Dietary tags that apply to this item, empty list if none"
    )


class MenuCategory(BaseModel):
    """A section of the menu (e.g., Appetizers, Mains)."""
    category_name: str = Field(description="Name of the menu section")
    items: list[MenuItem] = Field(description="Items in this section")


class ParsedMenu(BaseModel):
    """Structured representation of a restaurant menu."""
    restaurant_name: str = Field(description="Name of the restaurant")
    cuisine_type: str = Field(description="Type of cuisine (e.g., Italian, Japanese, American)")
    categories: list[MenuCategory] = Field(description="Menu sections with their items")
    price_range: Literal["$", "$$", "$$$", "$$$$"] = Field(
        description="Overall price range: $ (under 15), $$ (15-30), $$$ (30-60), $$$$ (60+)"
    )

menu_agent = Agent(
    get_model("claude-sonnet-4-6"),
    output_type=ParsedMenu,
    instructions="""## Role
You are a restaurant menu parser that extracts structured data from menu text.

## Task
Parse the provided menu text and extract all items with their details.

## Constraints
- If a price is not listed, set it to null
- Infer dietary tags only when explicitly stated or obvious from ingredients
  (e.g., "tofu stir-fry" is vegetarian, but don't guess)
- If the restaurant name isn't in the text, use "Unknown Restaurant"
- Categorize items into their menu sections; if no sections exist, use "General"
""",
)

menu_text_1 = """
SAKURA JAPANESE KITCHEN

🥢 STARTERS
Edamame (v) - Steamed soybeans with sea salt — $6
Gyoza - Pan-fried pork dumplings (6 pcs) — $9
Miso Soup (v, gf) - Traditional soybean soup with tofu and seaweed — $5

🍣 SUSHI ROLLS
California Roll - Crab, avocado, cucumber (8 pcs) — $12
Spicy Tuna Roll - Fresh tuna with spicy mayo, jalapeño (8 pcs) — $14
Dragon Roll - Shrimp tempura, eel, avocado (8 pcs) — $18

🍜 ENTREES
Chicken Teriyaki - Grilled chicken with teriyaki glaze, steamed rice — $16
Pad Thai - Rice noodles, shrimp, peanuts, bean sprouts (contains nuts) — $15
Vegetable Tempura (v) - Assorted vegetables in light batter — $13

🍡 DESSERTS
Mochi Ice Cream (gf) - Green tea, mango, or strawberry (3 pcs) — $7
Matcha Cheesecake — $9
"""

result = await menu_agent.run(menu_text_1)
menu = result.output

print(f"Restaurant: {menu.restaurant_name}")
print(f"Cuisine:    {menu.cuisine_type}")
print(f"Price Range: {menu.price_range}")
print(f"Categories: {len(menu.categories)}")
print()

for category in menu.categories:
    print(f"--- {category.category_name} ---")
    for item in category.items:
        tags = f" [{', '.join(item.dietary_tags)}]" if item.dietary_tags else ""
        price = f"${item.price:.2f}" if item.price else "no price"
        print(f"  {item.name}: {price}{tags}")
    print()

menu_text_2 = """
The Rustic Table - Farm to Fork American Cuisine

Appetizers
  Crispy Brussels Sprouts with balsamic glaze (v, gf) .... 11
  Loaded Nachos with pulled pork, cheddar, jalapeños .... 14
  Soup of the Day — ask your server .... 8

Burgers & Sandwiches
  Classic Smash Burger - double patty, American cheese, pickles .... 17
  BBQ Pulled Pork Sandwich - house-smoked pork, coleslaw .... 16
  Impossible Burger (v) - plant-based patty, all the fixings .... 18

Sides
  Truffle Fries .... 7
  Mac & Cheese .... 6
  Garden Salad (v, gf) .... 5

Desserts
  Bourbon Pecan Pie (contains nuts) .... 10
  Chocolate Lava Cake .... 12
"""

result_2 = await menu_agent.run(menu_text_2)
menu_2 = result_2.output

# Compare the two menus
print(f"{'':>20} {'Menu 1':>15} {'Menu 2':>15}")
print(f"{'Restaurant':>20} {menu.restaurant_name:>15} {menu_2.restaurant_name:>15}")
print(f"{'Cuisine':>20} {menu.cuisine_type:>15} {menu_2.cuisine_type:>15}")
print(f"{'Price Range':>20} {menu.price_range:>15} {menu_2.price_range:>15}")
print(f"{'Categories':>20} {len(menu.categories):>15} {len(menu_2.categories):>15}")

total_items_1 = sum(len(c.items) for c in menu.categories)
total_items_2 = sum(len(c.items) for c in menu_2.categories)
print(f"{'Total Items':>20} {total_items_1:>15} {total_items_2:>15}")

import spacy

nlp = spacy.load("en_core_web_sm")


def analyze_menu_text(text: str) -> str:
    """Analyze menu text with SpaCy to extract named entities and key noun phrases.

    Args:
        text: The full menu text to analyze.
    """
    doc = nlp(text)

    entities = [
        {"text": ent.text, "label": ent.label_}
        for ent in doc.ents
    ]

    # Extract noun chunks that might be ingredient or dish references
    noun_phrases = [chunk.text for chunk in doc.noun_chunks][:20]  # limit to top 20

    return (
        f"Named entities: {entities}\n"
        f"Key noun phrases: {noun_phrases}\n"
        f"Token count: {len(doc)}, Sentence count: {len(list(doc.sents))}"
    )


class EnrichedMenu(BaseModel):
    """Menu data enriched with NLP analysis."""
    restaurant_name: str = Field(description="Name of the restaurant")
    cuisine_type: str = Field(description="Type of cuisine")
    total_items: int = Field(description="Total number of menu items")
    avg_price: float | None = Field(description="Average price across all items, null if no prices listed")
    dietary_options: list[str] = Field(description="List of dietary accommodations available (e.g., vegetarian, gluten-free)")
    signature_dishes: list[str] = Field(description="2-3 dishes that seem most distinctive or premium")
    spacy_entities_summary: str = Field(description="Brief summary of what SpaCy found in the text (entities, key ingredients)")


enriched_agent = Agent(
    get_model("claude-sonnet-4-6"),
    output_type=EnrichedMenu,
    tools=[analyze_menu_text],
    instructions="""## Role
You are a restaurant data analyst that combines LLM understanding with NLP tool analysis.

## Task
Analyze the menu text to produce an enriched summary. Use the SpaCy tool to get
precise entity and noun phrase extraction, then combine that with your own
understanding of the menu content.

## Steps
1. First, use the analyze_menu_text tool on the full menu text
2. Then, synthesize the tool results with your own analysis
3. Return the enriched menu summary
""",
)

result = await enriched_agent.run(menu_text_1)
enriched = result.output

print(f"Restaurant:       {enriched.restaurant_name}")
print(f"Cuisine:          {enriched.cuisine_type}")
print(f"Total Items:      {enriched.total_items}")
print(f"Avg Price:        ${enriched.avg_price:.2f}" if enriched.avg_price else "Avg Price: N/A")
print(f"Dietary Options:  {', '.join(enriched.dietary_options)}")
print(f"Signature Dishes: {', '.join(enriched.signature_dishes)}")
print(f"\nSpaCy Analysis:   {enriched.spacy_entities_summary}")

def summarizer(
    max_sentences: int = 3,
    audience: str = "general",
    focus: str = "key information",
) -> str:
    """Summarization prompt template."""
    return (
        f"Summarize the text in {max_sentences} sentences or fewer. "
        f"Write for a {audience} audience. Focus on {focus}. "
        f"Do not start with 'This text discusses' or similar phrasing."
    )


def entity_extractor(
    entity_types: list[str],
    output_format: str = "list",
) -> str:
    """Entity extraction prompt template."""
    types_str = ", ".join(entity_types)
    return (
        f"Extract all {types_str} entities from the text. "
        f"Return them as a {output_format}. "
        f"Only include entities explicitly mentioned — do not infer or guess. "
        f"If no entities of a type are found, indicate that explicitly."
    )


def classifier(
    categories: list[str],
    description: str = "",
) -> str:
    """Text classification prompt template."""
    cats_str = ", ".join(categories)
    prompt = (
        f"Classify the text into exactly one of these categories: {cats_str}. "
        f"Respond with only the category name."
    )
    if description:
        prompt += f" Context: {description}"
    return prompt


# Test the library on a sample text
sample = (
    "Tesla CEO Elon Musk announced a new Gigafactory in Austin, Texas, "
    "expected to create 10,000 jobs and produce next-generation battery cells. "
    "The $5 billion investment reflects Tesla's push to dominate the EV market "
    "amid growing competition from BYD and Rivian."
)

# Template 1: Summarize
agent = Agent(get_model("claude-haiku-4-5"), instructions=summarizer(max_sentences=2, audience="investor"))
result = await agent.run(sample)
print(f"Summary: {result.output}\n")

# Template 2: Extract entities
agent = Agent(get_model("claude-haiku-4-5"), instructions=entity_extractor(["person", "organization", "location", "money"]))
result = await agent.run(sample)
print(f"Entities: {result.output}\n")

# Template 3: Classify
agent = Agent(get_model("claude-haiku-4-5"), instructions=classifier(
    ["technology", "finance", "politics", "sports", "science"],
    description="News article headlines and excerpts"
))
result = await agent.run(sample)
print(f"Category: {result.output}")

async def analyze_article(text: str) -> dict:
    """Multi-step analysis pipeline: classify → extract → summarize."""
    # Step 1: Classify the article
    classify_agent = Agent(
        get_model("claude-haiku-4-5"),
        instructions=classifier(
            ["technology", "finance", "politics", "sports", "science", "health"]
        ),
    )
    category_result = await classify_agent.run(text)
    category = category_result.output

    # Step 2: Extract entities (customize by category)
    if category.lower() in ["technology", "science"]:
        entity_types = ["person", "organization", "technology", "product"]
    elif category.lower() == "finance":
        entity_types = ["person", "organization", "money", "location"]
    else:
        entity_types = ["person", "organization", "location"]

    extract_agent = Agent(
        get_model("claude-haiku-4-5"),
        instructions=entity_extractor(entity_types),
    )
    entities_result = await extract_agent.run(text)

    # Step 3: Summarize with context
    summary_agent = Agent(
        get_model("claude-haiku-4-5"),
        instructions=summarizer(
            max_sentences=2,
            audience="analyst",
            focus=f"the key {category.lower()} implications",
        ),
    )
    summary_result = await summary_agent.run(text)

    return {
        "category": category,
        "entities": entities_result.output,
        "summary": summary_result.output,
    }


# Run the pipeline
analysis = await analyze_article(sample)

print(f"Category: {analysis['category']}")
print(f"\nEntities:\n{analysis['entities']}")
print(f"\nSummary: {analysis['summary']}")