Lab 2: Type-Safe Pydantic Tool Calling & Live Temporal Grounding

import os
import requests
from datetime import datetime
from dotenv import load_dotenv
from pydantic import BaseModel, Field, field_validator
import mlx.core as mx
from mlx_vlm import load, generate

# 1. Load environment variables
load_dotenv()
SERPER_API_KEY = os.getenv("SERPER_API_KEY")

# 2. Dynamic Temporal Grounding: Capture live system hardware date
CURRENT_SYSTEM_DATE = datetime.now().strftime("%B %d, %Y")
print(f"Temporal Grounding Anchor: {CURRENT_SYSTEM_DATE}")

# 3. Verify Apple Metal acceleration & Load Gemma 4
assert mx.metal.is_available(), "Apple Metal GPU acceleration required!"
MODEL_PATH = "AIModels/gemma-4-12B-it-qat-4bit"
model, processor = load(MODEL_PATH)
print("Gemma 4 (12B) successfully loaded into Apple Unified Memory.")

class TradeOrder(BaseModel):
    ticker: str = Field(..., pattern=r"^[A-Za-z]{1,5}$", description="Stock ticker symbol")
    shares: int = Field(..., gt=0, le=10000, description="Number of shares to execute")
    limit_price: float = Field(..., gt=0.0, description="Strictly positive limit price")

    @field_validator("ticker")
    @classmethod
    def clean_ticker(cls, v: str) -> str:
        return v.strip().upper()

# Test automated normalization on dirty LLM output
raw_llm_payload = {"ticker": "  nvda  ", "shares": "50", "limit_price": "142.50"}
order = TradeOrder.model_validate(raw_llm_payload)

print(f"Validated Ticker: {order.ticker}")      # 'NVDA' (coerced & uppercased)
print(f"Validated Shares: {order.shares}")      # 50 (int)
print(f"Validated Limit Price: ${order.limit_price:.2f}")  # 142.50 (float)

from typing import Literal, Optional

class LiveQuoteParams(BaseModel):
    ticker: str = Field(..., pattern=r"^[A-Za-z]{1,5}$", description="Stock ticker symbol")
    
    @field_validator("ticker")
    @classmethod
    def to_upper(cls, v: str) -> str:
        return v.strip().upper()

class ValuationModelParams(BaseModel):
    ticker: str = Field(..., pattern=r"^[A-Za-z]{1,5}$")
    current_price: float = Field(..., gt=0.0)
    earnings_per_share: float = Field(..., gt=0.0)
    industry_benchmark_pe: float = Field(default=25.0, gt=0.0, le=150.0)

class AgentToolDecision(BaseModel):
    reasoning: str = Field(..., description="Deductive rationale explaining the action")
    tool: Literal["fetch_live_quote", "evaluate_valuation_ratio", "none"]
    quote_params: Optional[LiveQuoteParams] = None
    valuation_params: Optional[ValuationModelParams] = None

# Export OpenAPI-compliant JSON schema for Gemma 4
DECISION_SCHEMA = AgentToolDecision.model_json_schema()

def fetch_live_quote(params: LiveQuoteParams) -> dict:
    url = "https://google.serper.dev/search"
    headers = {
        "X-API-KEY": SERPER_API_KEY,
        "Content-Type": "application/json"
    }
    payload = {"q": f"{params.ticker} stock price live today market", "num": 5}
    
    response = requests.post(url, headers=headers, json=payload, timeout=8)
    data = response.json()
    
    # Extract live snippet and stamp with live hardware date
    snippet = data.get("organic", [{}])[0].get("snippet", "No market data found.")
    return {
        "ticker": params.ticker,
        "telemetry": snippet,
        "as_of_date": CURRENT_SYSTEM_DATE  # Dynamic Temporal Anchor
    }

def evaluate_valuation_ratio(params: ValuationModelParams) -> dict:
    implied_pe = params.current_price / params.earnings_per_share
    fair_value = params.earnings_per_share * params.industry_benchmark_pe
    variance_pct = ((params.current_price - fair_value) / fair_value) * 100
    
    return {
        "ticker": params.ticker,
        "implied_pe": round(implied_pe, 2),
        "fair_value": round(fair_value, 2),
        "variance_pct": round(variance_pct, 2),
        "assessment": "Overvalued" if variance_pct > 10 else "Undervalued" if variance_pct < -10 else "Fair Value"
    }

import json

def run_agent(user_query: str) -> str:
    system_prompt = (
        f"You are an autonomous Wall Street research assistant. "
        f"The current date is {CURRENT_SYSTEM_DATE}. "
        f"You must respond strictly with valid JSON conforming to this schema:\n"
        f"{json.dumps(DECISION_SCHEMA)}"
    )
    
    # Phase 1: Structured Intent Analysis
    conversation = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_query}
    ]
    formatted = processor.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
    raw_decision = generate(model, processor, prompt=formatted, max_tokens=256, temp=0.1)
    
    # Phase 2: Schema Validation
    decision = AgentToolDecision.model_validate_json(raw_decision)
    
    # Phase 3: Tool Dispatch
    if decision.tool == "fetch_live_quote":
        obs = fetch_live_quote(decision.quote_params)
    elif decision.tool == "evaluate_valuation_ratio":
        obs = evaluate_valuation_ratio(decision.valuation_params)
    else:
        obs = {"info": "No external tool required."}
        
    # Phase 4: Executive Synthesis
    synthesis_prompt = (
        f"You are a Senior Equity Analyst. Today's verified date is {CURRENT_SYSTEM_DATE}.\n"
        f"User Query: {user_query}\n"
        f"Tool Observation: {json.dumps(obs)}\n"
        f"Provide a concise, professional research brief. Do NOT hallucinate past dates."
    )
    
    formatted_synthesis = processor.apply_chat_template(
        [{"role": "user", "content": synthesis_prompt}],
        tokenize=False,
        add_generation_prompt=True
    )
    return generate(model, processor, prompt=formatted_synthesis, max_tokens=384, temp=0.2)