Skip to main content

Theory of Mind modeling for Software Engineering assistants

Project description

iclr2026-tom-swe-intro-fig

TOM-SWE: User Mental Modeling For Software Engineering Agents

pytest pre-commit mypy PyPI

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}
}

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

tom_swe-1.0.2.tar.gz (40.6 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

tom_swe-1.0.2-py3-none-any.whl (50.6 kB view details)

Uploaded Python 3

File details

Details for the file tom_swe-1.0.2.tar.gz.

File metadata

  • Download URL: tom_swe-1.0.2.tar.gz
  • Upload date:
  • Size: 40.6 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for tom_swe-1.0.2.tar.gz
Algorithm Hash digest
SHA256 ab101cda50650b9132b79dc93d6f1f8637e3a773b9f1988caad8cdc17bc8150c
MD5 a23da950ea2d64c04b3da17bdd1b3b1c
BLAKE2b-256 d99b485a5787dcd0796ef0d1a04fe275ac74acd9114225f1529e684318987e02

See more details on using hashes here.

Provenance

The following attestation bundles were made for tom_swe-1.0.2.tar.gz:

Publisher: publish-to-pypi.yml on OpenHands/ToM-SWE

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tom_swe-1.0.2-py3-none-any.whl.

File metadata

  • Download URL: tom_swe-1.0.2-py3-none-any.whl
  • Upload date:
  • Size: 50.6 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for tom_swe-1.0.2-py3-none-any.whl
Algorithm Hash digest
SHA256 9490d96e13aeeb7da8857bed5b54ed9fae6fbc7f2b5e8065eaa15698d519ba6e
MD5 4ba12ee4361a50a28f49008d8e71f14a
BLAKE2b-256 3f4df15f7e9f20bdedc8f112eff1cebce5d50b852c568e66855382cb903fd4e9

See more details on using hashes here.

Provenance

The following attestation bundles were made for tom_swe-1.0.2-py3-none-any.whl:

Publisher: publish-to-pypi.yml on OpenHands/ToM-SWE

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page