Using Agent

WeasToolkit exposes LangChain-compatible tools that let an agent inspect and edit a WeasWidget. One can use the langchain-widget package to create a chat agent with the toolkit. First, install the required packages:

pip install langchain-widget[openai]

Then, you can create a toolkit instance by passing in a WeasWidget instance.

Important

The WeasToolkit requires an LLM that supports function calling, such as GPT. Make sure to set up your OpenAI API key in the environment variable OPENAI_API_KEY. Check the LangChain OpenAI documentation for more details.

Example

Run the following code to create a LangChain chat widget integrated with weas_widget. And try to ask:

  • “Load a Si diamond conventional cell and repeat 2x2x2”

  • “Select atom 0 and atom 1.”

  • “Replace the selected atoms with Ge”

  • “Summarize the structure”

from weas_widget import WeasWidget, WeasToolkit
from langchain_openai import ChatOpenAI
from langchain_widget import LangChainWidget
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
import ipywidgets as ipw

load_dotenv()

viewer = WeasWidget()

chat_model = ChatOpenAI(model="gpt-4o")
chat = LangChainWidget(
   chat_model=chat_model,
   tools=WeasToolkit(viewer=viewer),
   title="WEAS Agent Chat",
   system_prompt=(
      "You are a scientific assistant. "
      "Use the available tools to inspect and manipulate the 3D structure."
   ),
   sidebar_open=False,
)
ipw.HBox([viewer, chat])
Agent Tools Example

Extending the toolkit

You can add your own tools in two ways:

  1. Pass tools directly:

toolkit = WeasToolkit(viewer, extra_tools=[my_tool])
tools = toolkit.tools

  1. Register entry points under weas_widget.tools that return a tool or a tool factory.