app.py

import os
import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import random
from typing import List
from rcwa import Material, Layer, LayerStack, Source, Solver
from smolagents import tool, CodeAgent, InferenceClientModel, stream_to_gradio

# --- Constants ---
start_wl = 0.32
stop_wl = 0.80
step_wl = 0.01
wavelengths = np.arange(start_wl, stop_wl + step_wl, step_wl)
materials = ['Si', 'Si3N4', 'SiO2', 'AlN']

@tool
def simulate_spectrum_10nm(layer_order: List[str]) -> List[float]:
    """
    Simulates the optical transmission spectrum for a given sequence of material layers at 10nm thickness.

    Args:
        layer_order (List[str]): A list of material names (e.g., ["Si", "SiO2", "AlN"]) representing the order of layers in the optical stack.

    Returns:
        List[float]: The transmission spectrum across a predefined wavelength range.
    """
    source = Source(wavelength=start_wl)
    reflection_layer = Layer(n=1.0)
    transmission_layer = Layer(material=Material("Si"))
    try:
        layers = [Layer(material=Material(m), thickness=0.01) for m in layer_order]
        stack = LayerStack(*layers, incident_layer=reflection_layer, transmission_layer=transmission_layer)
        solver = Solver(stack, source, (1, 1))
        result = solver.solve(wavelength=wavelengths)
        return np.array(result['TTot']).tolist()
    except Exception as e:
        return []

@tool
def simulate_spectrum_100nm(layer_order: List[str]) -> List[float]:
    """
    Simulates the optical transmission spectrum for a given sequence of material layers at 100nm thickness.

    Args:
        layer_order (List[str]): A list of material names (e.g., ["Si", "SiO2", "AlN"]) representing the order of layers in the optical stack.

    Returns:
        List[float]: The transmission spectrum across a predefined wavelength range.
    """
    source = Source(wavelength=start_wl)
    reflection_layer = Layer(n=1.0)
    transmission_layer = Layer(material=Material("Si"))
    try:
        layers = [Layer(material=Material(m), thickness=0.1) for m in layer_order]
        stack = LayerStack(*layers, incident_layer=reflection_layer, transmission_layer=transmission_layer)
        solver = Solver(stack, source, (1, 1))
        result = solver.solve(wavelength=wavelengths)
        return np.array(result['TTot']).tolist()
    except Exception as e:
        return []

@tool
def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
    """
    Computes the cosine similarity between two vectors.

    Args:
        vec1 (List[float]): The first vector.
        vec2 (List[float]): The second vector.

    Returns:
        float: A similarity score between -1 and 1.
    """
    a, b = np.array(vec1), np.array(vec2)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

# --- Target Spectrum Generator ---
def get_target_spectrum(layer_order, thickness=0.1):
    source = Source(wavelength=start_wl)
    reflection_layer = Layer(n=1.0)
    transmission_layer = Layer(material=Material("Si"))
    try:
        layers = [Layer(material=Material(m), thickness=thickness) for m in layer_order]
        stack = LayerStack(*layers, incident_layer=reflection_layer, transmission_layer=transmission_layer)
        solver = Solver(stack, source, (1, 1))
        result = solver.solve(wavelength=wavelengths)
        return np.array(result['TTot']).tolist()
    except Exception:
        return None

# --- Model Setup ---
from smolagents import LiteLLMModel
openai_key = os.getenv("OPENAI_API_KEY")
model = LiteLLMModel(model_id="openai/gpt-4.1-mini", temperature=0, api_key=openai_key)

# --- Agent Setup ---
agent_10nm_simulator = CodeAgent(
    tools=[simulate_spectrum_10nm],
    model=model,
    stream_outputs=True,
    name="agent_10nm_simulator",
    description="You are an AI agent that uses tools to simulate optical spectra for materials with thickness 10nm. You must provide the simulated response back. Do not provide any other information. "
)

agent_10nm_simulator.prompt_templates['managed_agent'] = {
    "task": """You're an assistant agent named '{{name}}'.
You have been given this task:
---
{{task}}
---
Just return the result of your tool call. Do not add explanations or formatting.
Call a tool immediately and use `final_answer(...)` to return the result.
""",
    "report": """{{final_answer}}"""  # Minimal required key
}



agent_100nm_simulator = CodeAgent(
    tools=[simulate_spectrum_100nm],
    model=model,
    stream_outputs=True,
    name="agent_100nm_simulator",
    description="You are an AI agent that uses tools to simulate optical spectra for materials with thickness 100nm. You must provide the simulated response back. Do not provide any other information."
)


agent_100nm_simulator.prompt_templates['managed_agent'] = {
    "task": """You're an assistant agent named '{{name}}'.
You have been given this task:
---
{{task}}
---
Just return the result of your tool call. Do not add explanations or formatting.
Call a tool immediately and use `final_answer(...)` to return the result.
""",
    "report": """{{final_answer}}"""  # Minimal required key
}


coordinator = CodeAgent(
    tools=[cosine_similarity],
    managed_agents=[agent_10nm_simulator, agent_100nm_simulator],
    model=model,
    stream_outputs=True,
    additional_authorized_imports = ["numpy"]
)

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown(
        """
    # 🧠 Multi-Agent Thin-Film Stack Optimizer (10nm + 100nm RCWA)
    
    This interactive demo showcases a **multi-agent AI system** that cooperatively solves an inverse optics problem using physics-based simulation.  
    Instead of relying on a single agent or a single simulator, the **Coordinator Agent** orchestrates two specialized agents, each tuned to simulate optical spectra for different layer thicknesses (10nm and 100nm), and uses a comparison tool to evaluate results.
    
    ---
    
    ### 🤖 Objective: Discover the correct layer **order** and **thickness**
    
    Given only a target transmission spectrum, the multi-agent system must identify:
    - The correct **material permutation** (from a fixed set), and
    - The correct **thickness choice** (10nm or 100nm, uniformly applied)
    
    **Constraints:**
    - Materials: `Si`, `Si₃N₄`, `SiO₂`, `AlN` (used once each)
    - Layer thickness options: `10nm` or `100nm`
    - Terminate when `cosine_similarity > 0.999`
    
    ---
    
    ### 🧬 Agent Architecture Overview
    
    - 🧠 **Coordinator Agent** (`CodeAgent | gpt-4.1-mini`):  
      Receives the target spectrum and is responsible for reasoning, selecting candidates, and delegating simulations to sub-agents.
    
    - 🔧 **agent_10nm_simulator**:  
      Can simulate any material order using **10nm** thick layers via RCWA.
    
    - 🔧 **agent_100nm_simulator**:  
      Can simulate any material order using **100nm** thick layers via RCWA.
    
    - 📏 **cosine_similarity tool**:  
      Measures how close a simulated spectrum is to the target.
    
    ---
    
    ### 🔍 What’s Happening Under the Hood
    
    1. A **random 4-layer stack** is selected from `Si`, `Si₃N₄`, `SiO₂`, `AlN`, and simulated (at 100nm) to serve as the **target spectrum**.
    2. The **Coordinator Agent** is provided the target and access to two sub-agents:
        - `agent_10nm_simulator` for 10nm stacks
        - `agent_100nm_simulator` for 100nm stacks
    3. It explores permutations + thickness combinations by:
        - Calling a simulation agent
        - Comparing simulated output with the target using `cosine_similarity`
        - Continuing exploration until similarity exceeds threshold
    4. The system **halts automatically** once the optimal stack is found, and reports:
        - The matched material order
        - Thickness value
        - Number of permutations tried
    
    ---
    ### 📊 Visualization

        """
    )
    gr.Markdown("### 📊 Transmission Spectra of Layer Stack Designs for all 24 permutations of material order for thickness 100nm")
    gr.Image(value="121_resized.png", interactive=False)

    gr.Markdown("### 📊 Transmission Spectra of Layer Stack Designs for all 24 permutations of material order for thickness 10nm")
    gr.Image(value="122_resized.png", interactive=False)
    
    gr.Markdown("""
    ### 🤖 Multi-Agent System Overview
    
    #### 🧠 Coordinator Agent: `CodeAgent | openai/gpt-4.1-mini`
    
    - ✅ **Authorized Imports**: `['numpy']`
    
    ##### 🛠️ Tools:
    | Tool               | Description                                | Arguments                                 |
    |--------------------|--------------------------------------------|-------------------------------------------|
    | `cosine_similarity`| Computes the cosine similarity between two vectors | `vec1`, `vec2`: Lists of floats   |
    
    #### 🔧 Managed Agent: `agent_10nm_simulator | CodeAgent | openai/gpt-4.1-mini`
    
    - ✅ **Authorized Imports**: `[]`
    - 📝 **Description**: Simulates optical spectra for **10nm** thickness.
    
    ##### 🛠️ Tools:
    | Tool                  | Description                                  | Arguments                                  |
    |-----------------------|----------------------------------------------|--------------------------------------------|
    | `simulate_spectrum_10nm` | Simulates spectrum for 10nm layers.         | `layer_order`: List of materials           |
    
    ---
    
    #### 🔧 Managed Agent: `agent_100nm_simulator | CodeAgent | openai/gpt-4.1-mini`
    
    - ✅ **Authorized Imports**: `[]`
    - 📝 **Description**: Simulates optical spectra for **100nm** thickness.
    
    ##### 🛠️ Tools:
    | Tool                   | Description                                   | Arguments                                  |
    |------------------------|-----------------------------------------------|--------------------------------------------|
    | `simulate_spectrum_100nm` | Simulates spectrum for 100nm layers.          | `layer_order`: List of materials           |
    """)

    run_btn = gr.Button("🔁 Run Agent on Random Stack")
    true_order = gr.Textbox(label="True Material Order")
    prompt_box = gr.Textbox(label="Agent Prompt")
    chatbot = gr.Chatbot(label="Agent Reasoning Stream")

    def run_agent_streaming():
        true_order_val = random.sample(materials, 4)
        target_val = get_target_spectrum(true_order_val)
        true_order_display = ", ".join(true_order_val)

        if target_val is None:
            yield gr.update(value="Simulation failed"), gr.update(), gr.update()
            return

        prompt = f"""
        You are the Coordinator Agent. Your objective is to identify a 4-layer material stack **order** and **thickness** that reproduces a given target optical transmission spectrum.
        
        Constraints:
        - Materials: [Si, Si3N4, SiO2, AlN] (use each exactly once)
        - Two fixed thickness options for all layers: 10nm and 100nm
        
        You have access to the following:
        - agent_10nm_simulator: An agent that simulates a spectrum for a given material order with **10nm** layer thickness
        - agent_100nm_simulator: An agent that simulates a spectrum for a given material order with **100nm** layer thickness
        - cosine_similarity: Compares a predicted spectrum to the target spectrum
        
        Your task:
        1. Choose candidate layer orders and thickness options
        2. Call the appropriate agent to simulate the spectrum
        3. Use cosine_similarity to compare with the target
        4. Stop when similarity exceeds 0.999
        5. Report the matching order, thickness, and number of attempts
        
        Begin.
        
        Target spectrum: {target_val}
        """
        chat_history = []
        yield gr.update(value=true_order_display), gr.update(value=prompt), gr.update(value=[])

        for msg in stream_to_gradio(coordinator, task=prompt):
            if isinstance(msg, gr.ChatMessage):
                chat_history.append(("", msg.content))
            elif isinstance(msg, str):
                if chat_history:
                    chat_history[-1] = ("", msg)
                else:
                    chat_history.append(("", msg))
            yield gr.update(), gr.update(), gr.update(value=chat_history)

    run_btn.click(fn=run_agent_streaming, inputs=[], outputs=[true_order, prompt_box, chatbot])

    demo.launch(debug=True)