TOM-SWE: User Mental Modeling For Software Engineering Agents

News

• [10/2025] Beta testing program launched - join us in testing ToM-enhanced agents!

Introduction

ToM-SWE is a Theory of Mind package designed to enhance Software Engineering agents with personalized user understanding and adaptive behavior.

ToM-SWE integrates seamlessly with OpenHands and other SWE agent frameworks, providing consultation capabilities that help agents understand user intent, preferences, and working styles for improved task performance.

Join the Beta Test

pip install uv

uvx --python 3.12 --from git+https://github.com/XuhuiZhou/OpenHands@feature/tom-codeact-agent openhands

For details, please refer to the Google doc

Get Started

Install from PyPI

pip install tom-swe

PyPI page: https://pypi.org/project/tom-swe/

Install Locally

We recommend using a virtual environment with uv:

pip install uv
uv sync

[!NOTE] You can use any other package manager to install dependencies (e.g. pip, conda), but we strongly recommend using uv for the best development experience.

Set up LLM API Credentials

Create a .env file with your credentials:

LITELLM_API_KEY=your_api_key_here
LITELLM_BASE_URL=your_proxy_endpoint
DEFAULT_LLM_MODEL=litellm_proxy/claude-sonnet-4-20250514

Easy Sample Demo

You can test the consultation functionality with a simple example:

from tom_swe.tom_module import TomModule
import asyncio

async def demo():
    tom = TomModule()

    # Consult on user preferences
    consultation = await tom.consult(
        user_id="demo_user",
        current_context="User wants to implement a new feature"
    )
    print(consultation)

asyncio.run(demo())

Or run the included example:

uv run python example.py           # See consultation functionality in action
uv run tom-config                  # Interactive LLM setup

Core Features

• Three-Tier Memory: Cleaned sessions → Session analyses → User profiles
• Agent Consultation: Provides personalized guidance and recommendations for SWE agents
• User Behavior Analysis: LLM-powered psychological insights and preferences
• OpenHands Integration: Use TomCodeActAgent for automatic instruction enhancement

Main Commands

# User analysis
uv run user-analysis --user-id <user_id>
uv run user-analysis --all-users --sample-size 100

# Theory of Mind analysis
uv run tom-test                      # Test on sample users
uv run tom-analyze                   # Full analysis

# RAG document analysis
uv run rag-agent

OpenHands Integration

Configure OpenHands to use Tom-enhanced agent:

default_agent = "TomCodeActAgent"

The agent automatically:

1. Provides consultation and personalized guidance to SWE agents
2. Processes user sessions for better understanding
3. Shows progress during analysis

Prompts

The prompts are stored in tom_swe/prompts/registry.py. You can also find some prompts in tom_swe/generation/dataclass.py

Requirements

• Python 3.8+
• uv package manager
• LLM API key (contact All Hands AI for access)

Citation

If you use ToM-SWE in your research, please cite:

@software{zhou2024tomswe,
  title = {TOM-SWE: User Mental Modeling For Software Engineering Agents},
  author = {Xuhui Zhou and Valerie Chen and Zora Zhiruo Wang and
Graham Neubig and Maarten Sap and Xingyao Wang},
  year = {2025},
  url = {https://github.com/All-Hands-AI/TOM-SWE}
}