Create app.py

app.py

@@ -0,0 +1,216 @@
+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_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
+    This demo simulates an AI agent coordinating spectrum simulations at 10nm and 100nm thickness to match a target.
+    """)
+    gr.Markdown("""
+    CodeAgent | openai/gpt-4.1-mini  
+├── ✅ Authorized imports: ['numpy']  
+├── 🛠️ Tools:  
+│   ┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓  
+│   ┃ Name              ┃ Description                                ┃ Arguments                                  ┃  
+│   ┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩  
+│   │ cosine_similarity │ Computes the cosine similarity between two │ vec1 (`array`): The first vector.          │  
+│   │                   │ vectors.                                   │ vec2 (`array`): The second vector.         │  
+│   └───────────────────┴────────────────────────────────────────────┴────────────────────────────────────────────┘  
+└── 🤖 Managed agents:  
+    ├── agent_10nm_simulator | CodeAgent | openai/gpt-4.1-mini  
+    │   ├── ✅ Authorized imports: []  
+    │   ├── 📝 Description: Simulates optical spectra for 10nm thickness.  
+    │   └── 🛠️ Tools:  
+    │       ┏━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓  
+    │       ┃ Name                   ┃ Description                         ┃ Arguments                            ┃  
+    │       ┡━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩  
+    │       │ simulate_spectrum_10nm │ Simulates spectrum for 10nm layers. │ layer_order (`array`): List of materials │  
+    │       └────────────────────────┴─────────────────────────────────────┴──────────────────────────────────────┘  
+    └── agent_100nm_simulator | CodeAgent | openai/gpt-4.1-mini  
+        ├── ✅ Authorized imports: []  
+        ├── 📝 Description: Simulates optical spectra for 100nm thickness.  
+        └── 🛠️ Tools:  
+            ┏━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓  
+            ┃ Name                    ┃ Description                         ┃ Arguments                           ┃  
+            ┡━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩  
+            │ simulate_spectrum_100nm │ Simulates spectrum for 100nm layers.│ layer_order (`array`): 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)