chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,53 @@
|
||||
name: Bug Report
|
||||
description: File a bug report to help us improve Trae Agent
|
||||
title: "[Bug]: "
|
||||
labels: ["type/bug", "status/need_triage"]
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Thanks for taking the time to fill out this bug report! Before reporting a bug, please make sure you have searched https://github.com/bytedance/trae-agent/issues to see if there are any existing issues that cover the same problem.
|
||||
|
||||
- type: textarea
|
||||
id: what-happened
|
||||
attributes:
|
||||
label: What happened?
|
||||
description: Please provide a clear and concise description of what the bug is. Note that please don't upload your API secrate token to the bug report.
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: what-expected
|
||||
attributes:
|
||||
label: What did you expect to happen?
|
||||
description: Please provide a clear and concise description of what you expected to happen.
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: traceback
|
||||
attributes:
|
||||
label: Traceback
|
||||
description: Please provide the traceback if an exception occurs.
|
||||
validations:
|
||||
required: false
|
||||
|
||||
- type: textarea
|
||||
id: env-info
|
||||
attributes:
|
||||
label: What is your system, Python, dependency version?
|
||||
description: Please provide your system, Python, dependency version.
|
||||
placeholder: |
|
||||
- OS: [e.g. Ubuntu 20.04]
|
||||
- Python: [e.g. Python 3.10]
|
||||
- Dependency Version: [e.g. transformers 4.32.1]
|
||||
validations:
|
||||
required: false
|
||||
|
||||
- type: textarea
|
||||
id: additional-info
|
||||
attributes:
|
||||
label: Additional information that you believe is relevant to this bug
|
||||
description: Add any other context about the problem here.
|
||||
validations:
|
||||
required: false
|
||||
@@ -0,0 +1,8 @@
|
||||
blank_issues_enabled: false
|
||||
contact_links:
|
||||
- name: Trae Agent Discussions
|
||||
url: https://github.com/bytedance/trae-agent/discussions
|
||||
about: For general questions, roadmap, and ideas, please discuss here.
|
||||
- name: Trae AI IDE Community
|
||||
url: https://discord.gg/VwaQ4ZBHvC
|
||||
about: For all inquiries related to the product, please join the Discord community.
|
||||
@@ -0,0 +1,32 @@
|
||||
name: Feature Request
|
||||
description: Suggest a new feature or feature update for this project
|
||||
labels: ['type/feature', 'status/need-triage']
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Thanks for taking the time to fill out this feature request form! Before submitting a feature request, please make sure you have searched https://github.com/bytedance/trae-agent/issues to see if there are any existing issues that cover the same idea.
|
||||
|
||||
- type: textarea
|
||||
id: description
|
||||
attributes:
|
||||
label: What feature would you like to be added or updated?
|
||||
description: A clear and concise description of the feature request.
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: reason
|
||||
attributes:
|
||||
label: Why do you need this feature?
|
||||
description: A clear and concise description of the reason why this feature is needed.
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: additional-info
|
||||
attributes:
|
||||
label: Additional information that you believe is relevant to this feature request
|
||||
description: Add any other context about the idea here.
|
||||
validations:
|
||||
required: false
|
||||
@@ -0,0 +1,32 @@
|
||||
name: Feature Proposal
|
||||
description: Propose a new feature or enhancement for the trae-agent project
|
||||
labels: ['type/feature', 'status/need-triage']
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Thank you for contributing your ideas! Please check existing issues at https://github.com/bytedance/trae-agent/issues to avoid duplicates before submitting your proposal.
|
||||
|
||||
- type: textarea
|
||||
id: feature
|
||||
attributes:
|
||||
label: Describe the feature you want to propose
|
||||
description: Provide a detailed explanation of the feature or improvement you suggest.
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: motivation
|
||||
attributes:
|
||||
label: What problem does this feature solve or what benefit does it bring?
|
||||
description: Explain why this feature is important or how it will improve the project.
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: implementation-details
|
||||
attributes:
|
||||
label: Implementation details or suggestions (optional)
|
||||
description: Share any ideas or approaches for how this feature might be implemented.
|
||||
validations:
|
||||
required: false
|
||||
@@ -0,0 +1,24 @@
|
||||
name: Question
|
||||
description: Ask a question about Trae Agent
|
||||
labels: ['type/question', 'status/need-triage']
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Thanks for taking the time to fill out this question form! Before asking a question, please make sure you have searched https://github.com/bytedance/trae-agent/issues to see if there are any existing issues that cover the same question.
|
||||
|
||||
- type: textarea
|
||||
id: description
|
||||
attributes:
|
||||
label: What is your question?
|
||||
description: A clear and concise description of the question.
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: additional-info
|
||||
attributes:
|
||||
label: Additional information that you believe is relevant to this question
|
||||
description: Add any other context about the question here.
|
||||
validations:
|
||||
required: false
|
||||
@@ -0,0 +1,29 @@
|
||||
## Description
|
||||
|
||||
<!-- Add a brief description about this pull request including what it does, why it is needed, and other important information for the reviewers -->
|
||||
|
||||
## More Information
|
||||
|
||||
<!-- Add more in-depth information about this pull request, such as the changes made, the reasoning behind them, and any potential impacts. -->
|
||||
|
||||
## Validation
|
||||
|
||||
<!-- Introduce how to test this pull request. -->
|
||||
|
||||
## Linked Issues
|
||||
|
||||
<!--
|
||||
Link to any related issues or bugs.
|
||||
|
||||
**If this PR fully resolves the issue, use one of the following keywords to automatically close the issue when this PR is merged:**
|
||||
|
||||
- Closes #<issue_number>
|
||||
- Fixes #<issue_number>
|
||||
- Resolves #<issue_number>
|
||||
|
||||
*Example: `Resolves #123`*
|
||||
|
||||
**If this PR is only related to an issue or is a partial fix, simply reference the issue number without a keyword:**
|
||||
|
||||
*Example: `This PR makes progress on #456` or `Related to #789`*
|
||||
-->
|
||||
@@ -0,0 +1,39 @@
|
||||
name: Pre-commit
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
|
||||
jobs:
|
||||
pre-commit:
|
||||
|
||||
if: github.repository == 'bytedance/trae-agent'
|
||||
runs-on: ubuntu-latest
|
||||
name: Pre-commit checks
|
||||
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.12'
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v6
|
||||
|
||||
- name: Create virtual environment and install dependencies
|
||||
run: |
|
||||
make uv-sync
|
||||
|
||||
- name: Run pre-commit hooks
|
||||
run: |
|
||||
source .venv/bin/activate
|
||||
make uv-pre-commit
|
||||
@@ -0,0 +1,36 @@
|
||||
name: Unit Tests
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
|
||||
jobs:
|
||||
test:
|
||||
if: github.repository == 'bytedance/trae-agent'
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.12'
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v6
|
||||
|
||||
- name: Create virtual environment and install dependencies
|
||||
run: |
|
||||
make uv-sync
|
||||
|
||||
- name: Run unit tests
|
||||
run: |
|
||||
make uv-test
|
||||
+188
@@ -0,0 +1,188 @@
|
||||
# Python-generated files
|
||||
__pycache__/
|
||||
*.py[oc]
|
||||
build/
|
||||
dist/
|
||||
wheels/
|
||||
*.egg-info
|
||||
|
||||
# Virtual environments
|
||||
.venv
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# Node stuff:
|
||||
.node_modules/
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
#pdm.lock
|
||||
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||
# in version control.
|
||||
# https://pdm.fming.dev/#use-with-ide
|
||||
.pdm.toml
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be added to the global gitignore or merged into this project gitignore. For a PyCharm
|
||||
# project, it is recommended to uncomment the following lines to ignore the cache
|
||||
# files for the tool.
|
||||
#.idea/
|
||||
|
||||
trae-config-local.json
|
||||
trae_config.json
|
||||
trae_config.yaml
|
||||
|
||||
# Trajectories
|
||||
/trajectories/
|
||||
|
||||
# VS Code settings
|
||||
.vscode/
|
||||
!.vscode/launch.template.json
|
||||
|
||||
# Patch selection python binary
|
||||
py312/
|
||||
@@ -0,0 +1,34 @@
|
||||
repos:
|
||||
- repo: https://github.com/pre-commit/pre-commit-hooks
|
||||
rev: v5.0.0
|
||||
hooks:
|
||||
- id: trailing-whitespace
|
||||
- id: end-of-file-fixer
|
||||
- id: check-yaml
|
||||
- id: check-toml
|
||||
- id: check-added-large-files
|
||||
- id: detect-private-key
|
||||
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.12.1
|
||||
hooks:
|
||||
- id: ruff
|
||||
args: [ --fix ]
|
||||
- id: ruff-format
|
||||
|
||||
- repo: https://github.com/codespell-project/codespell
|
||||
rev: v2.4.1
|
||||
hooks:
|
||||
- id: codespell
|
||||
exclude: >
|
||||
(?x)^(
|
||||
.*\.jsonl
|
||||
)$
|
||||
|
||||
- repo: https://github.com/pre-commit/mirrors-mypy
|
||||
rev: v1.16.1
|
||||
hooks:
|
||||
- id: mypy
|
||||
exclude: ^(evaluation/patch_selection)
|
||||
additional_dependencies:
|
||||
- types-PyYAML
|
||||
@@ -0,0 +1 @@
|
||||
3.12
|
||||
+106
@@ -0,0 +1,106 @@
|
||||
# Contributing to Trae Agent
|
||||
|
||||
Thank you for your interest in contributing to Trae Agent! We welcome contributions of all kinds from the community.
|
||||
|
||||
## Ways to Contribute
|
||||
|
||||
There are many ways you can contribute to Trae Agent:
|
||||
|
||||
- **Code Contributions**: Add new features, fix bugs, or improve performance
|
||||
- **Documentation**: Improve README, add code comments, or create examples
|
||||
- **Bug Reports**: Submit detailed bug reports through issues
|
||||
- **Feature Requests**: Suggest new features or improvements
|
||||
- **Code Reviews**: Review pull requests from other contributors
|
||||
- **Community Support**: Help others in discussions and issues
|
||||
|
||||
## Development Setup
|
||||
|
||||
1. Fork the repository
|
||||
2. Clone your fork:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/bytedance/trae-agent.git
|
||||
cd trae-agent
|
||||
```
|
||||
|
||||
3. Set up your development environment:
|
||||
|
||||
```bash
|
||||
make install-dev
|
||||
make pre-commit-install
|
||||
```
|
||||
|
||||
## Running Tests
|
||||
|
||||
```bash
|
||||
make test
|
||||
```
|
||||
|
||||
## Development Process
|
||||
|
||||
1. Create a new branch:
|
||||
|
||||
```bash
|
||||
git checkout -b feature/amazing-feature
|
||||
```
|
||||
|
||||
2. Make your changes following our coding standards:
|
||||
- Write clear, documented code
|
||||
- Follow PEP 8 style guidelines
|
||||
- Add tests for new features
|
||||
- Update documentation as needed
|
||||
- Maintain type hints and add type checking when possible
|
||||
|
||||
3. Commit your changes:
|
||||
|
||||
```bash
|
||||
git commit -m 'Add some amazing feature'
|
||||
```
|
||||
|
||||
4. Push to your fork:
|
||||
|
||||
```bash
|
||||
git push origin feature/amazing-feature
|
||||
```
|
||||
|
||||
5. Open a Pull Request
|
||||
|
||||
## Pull Request Guidelines
|
||||
|
||||
- Fill in the pull request template completely
|
||||
- Include tests for new features
|
||||
- Update documentation as needed
|
||||
- Ensure all tests pass and there are no linting errors
|
||||
- Keep pull requests focused on a single feature or fix
|
||||
- Reference any related issues
|
||||
|
||||
## Code Style
|
||||
|
||||
- Follow PEP 8 guidelines
|
||||
- Use type hints where possible
|
||||
- Write descriptive docstrings
|
||||
- Keep functions and methods focused and single-purpose
|
||||
- Comment complex logic
|
||||
- Python version requirement: >= 3.12
|
||||
|
||||
## Community Guidelines
|
||||
|
||||
- Be respectful and inclusive
|
||||
- Follow our code of conduct
|
||||
- Help others learn and grow
|
||||
- Give constructive feedback
|
||||
- Stay focused on improving the project
|
||||
|
||||
## Need Help?
|
||||
|
||||
If you need help with anything:
|
||||
|
||||
- Check existing issues and discussions
|
||||
- Join our community channels
|
||||
- Ask questions in discussions
|
||||
|
||||
## License
|
||||
|
||||
By contributing to Trae Agent, you agree that your contributions will be licensed under the MIT License.
|
||||
|
||||
We appreciate your contributions to making Trae Agent better!
|
||||
@@ -0,0 +1,7 @@
|
||||
Copyright 2025 ByteDance Ltd. and/or its affiliates
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
||||
@@ -0,0 +1,57 @@
|
||||
.PHONY: help uv-venv uv-sync install-dev uv-pre-commit uv-test test pre-commit fix-format pre-commit-install pre-commit-run clean
|
||||
|
||||
# Default target
|
||||
help:
|
||||
@echo "Available commands:"
|
||||
@echo " install-dev - Create venv and install all dependencies (recommended for development)"
|
||||
@echo " uv-venv - Create a Python virtual environment using uv"
|
||||
@echo " uv-sync - Install all dependencies (including test/evaluation) using uv"
|
||||
@echo " uv-test - Run all tests (via uv, skips some external service tests)"
|
||||
@echo " test - Run all tests (skips some external service tests)"
|
||||
@echo " uv-pre-commit - Run pre-commit hooks on all files (via uv)"
|
||||
@echo " pre-commit-install- Install pre-commit hooks"
|
||||
@echo " pre-commit-run - Run pre-commit hooks on all files"
|
||||
@echo " pre-commit - Install and run pre-commit hooks on all files"
|
||||
@echo " fix-format - Fix formatting errors"
|
||||
@echo " clean - Clean up build artifacts and cache"
|
||||
|
||||
# Installation commands
|
||||
uv-venv:
|
||||
uv venv
|
||||
uv-sync:
|
||||
uv sync --all-extras
|
||||
install-dev: uv-venv uv-sync
|
||||
|
||||
# Pre-commit commands
|
||||
uv-pre-commit:
|
||||
uv run pre-commit run --all-files
|
||||
|
||||
pre-commit-install:
|
||||
pre-commit install
|
||||
pre-commit-run:
|
||||
pre-commit run --all-files
|
||||
pre-commit: pre-commit-install pre-commit-run
|
||||
|
||||
# fix formatting error
|
||||
fix-format:
|
||||
ruff format .
|
||||
ruff check --fix .
|
||||
|
||||
# Testing commands
|
||||
uv-test:
|
||||
SKIP_OLLAMA_TEST=true SKIP_OPENROUTER_TEST=true SKIP_GOOGLE_TEST=true uv run pytest tests/ -v --tb=short --continue-on-collection-errors
|
||||
test:
|
||||
SKIP_OLLAMA_TEST=true SKIP_OPENROUTER_TEST=true SKIP_GOOGLE_TEST=true uv run pytest
|
||||
|
||||
# Clean up
|
||||
clean:
|
||||
rm -rf build/
|
||||
rm -rf dist/
|
||||
rm -rf *.egg-info/
|
||||
rm -rf .pytest_cache/
|
||||
rm -rf .coverage
|
||||
rm -rf htmlcov/
|
||||
rm -rf .mypy_cache/
|
||||
rm -rf .ruff_cache/
|
||||
find . -type d -name __pycache__ -exec rm -rf {} +
|
||||
find . -name "*.pyc" -delete
|
||||
@@ -0,0 +1,286 @@
|
||||
# Trae Agent
|
||||
|
||||
[](https://arxiv.org/abs/2507.23370)
|
||||
[](https://www.python.org/downloads/) [](https://opensource.org/licenses/MIT)
|
||||
[](https://github.com/bytedance/trae-agent/actions/workflows/pre-commit.yml)
|
||||
[](https://github.com/bytedance/trae-agent/actions/workflows/unit-test.yml)
|
||||
[](https://discord.gg/VwaQ4ZBHvC)
|
||||
|
||||
**Trae Agent** is an LLM-based agent for general purpose software engineering tasks. It provides a powerful CLI interface that can understand natural language instructions and execute complex software engineering workflows using various tools and LLM providers.
|
||||
|
||||
For technical details please refer to [our technical report](https://arxiv.org/abs/2507.23370).
|
||||
|
||||
**Project Status:** The project is still being actively developed. Please refer to [docs/roadmap.md](docs/roadmap.md) and [CONTRIBUTING](CONTRIBUTING.md) if you are willing to help us improve Trae Agent.
|
||||
|
||||
**Difference with Other CLI Agents:** Trae Agent offers a transparent, modular architecture that researchers and developers can easily modify, extend, and analyze, making it an ideal platform for **studying AI agent architectures, conducting ablation studies, and developing novel agent capabilities**. This **_research-friendly design_** enables the academic and open-source communities to contribute to and build upon the foundational agent framework, fostering innovation in the rapidly evolving field of AI agents.
|
||||
|
||||
## ✨ Features
|
||||
|
||||
- 🌊 **Lakeview**: Provides short and concise summarisation for agent steps
|
||||
- 🤖 **Multi-LLM Support**: Works with OpenAI, Anthropic, Doubao, Azure, OpenRouter, Ollama and Google Gemini APIs
|
||||
- 🛠️ **Rich Tool Ecosystem**: File editing, bash execution, sequential thinking, and more
|
||||
- 🎯 **Interactive Mode**: Conversational interface for iterative development
|
||||
- 📊 **Trajectory Recording**: Detailed logging of all agent actions for debugging and analysis
|
||||
- ⚙️ **Flexible Configuration**: YAML-based configuration with environment variable support
|
||||
- 🚀 **Easy Installation**: Simple pip-based installation
|
||||
|
||||
## 🚀 Installation
|
||||
|
||||
### Requirements
|
||||
- UV (https://docs.astral.sh/uv/)
|
||||
- API key for your chosen provider (OpenAI, Anthropic, Google Gemini, OpenRouter, etc.)
|
||||
|
||||
### Setup
|
||||
|
||||
```bash
|
||||
git clone https://github.com/bytedance/trae-agent.git
|
||||
cd trae-agent
|
||||
uv sync --all-extras
|
||||
source .venv/bin/activate
|
||||
```
|
||||
|
||||
## ⚙️ Configuration
|
||||
|
||||
### YAML Configuration (Recommended)
|
||||
|
||||
1. Copy the example configuration file:
|
||||
```bash
|
||||
cp trae_config.yaml.example trae_config.yaml
|
||||
```
|
||||
|
||||
2. Edit `trae_config.yaml` with your API credentials and preferences:
|
||||
|
||||
```yaml
|
||||
agents:
|
||||
trae_agent:
|
||||
enable_lakeview: true
|
||||
model: trae_agent_model # the model configuration name for Trae Agent
|
||||
max_steps: 200 # max number of agent steps
|
||||
tools: # tools used with Trae Agent
|
||||
- bash
|
||||
- str_replace_based_edit_tool
|
||||
- sequentialthinking
|
||||
- task_done
|
||||
|
||||
model_providers: # model providers configuration
|
||||
anthropic:
|
||||
api_key: your_anthropic_api_key
|
||||
provider: anthropic
|
||||
openai:
|
||||
api_key: your_openai_api_key
|
||||
provider: openai
|
||||
|
||||
models:
|
||||
trae_agent_model:
|
||||
model_provider: anthropic
|
||||
model: claude-sonnet-4-20250514
|
||||
max_tokens: 4096
|
||||
temperature: 0.5
|
||||
```
|
||||
|
||||
**Note:** The `trae_config.yaml` file is ignored by git to protect your API keys.
|
||||
|
||||
### Using Base URL
|
||||
In some cases, we need to use a custom URL for the api. Just add the `base_url` field after `provider`, take the following config as an example:
|
||||
|
||||
```
|
||||
openai:
|
||||
api_key: your_openrouter_api_key
|
||||
provider: openai
|
||||
base_url: https://openrouter.ai/api/v1
|
||||
```
|
||||
**Note:** For field formatting, use spaces only. Tabs (\t) are not allowed.
|
||||
|
||||
### Environment Variables (Alternative)
|
||||
|
||||
You can also configure API keys using environment variables and store them in the .env file:
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY="your-openai-api-key"
|
||||
export OPENAI_BASE_URL="your-openai-base-url"
|
||||
export ANTHROPIC_API_KEY="your-anthropic-api-key"
|
||||
export ANTHROPIC_BASE_URL="your-anthropic-base-url"
|
||||
export GOOGLE_API_KEY="your-google-api-key"
|
||||
export GOOGLE_BASE_URL="your-google-base-url"
|
||||
export OPENROUTER_API_KEY="your-openrouter-api-key"
|
||||
export OPENROUTER_BASE_URL="https://openrouter.ai/api/v1"
|
||||
export DOUBAO_API_KEY="your-doubao-api-key"
|
||||
export DOUBAO_BASE_URL="https://ark.cn-beijing.volces.com/api/v3/"
|
||||
```
|
||||
|
||||
### MCP Services (Optional)
|
||||
|
||||
To enable Model Context Protocol (MCP) services, add an `mcp_servers` section to your configuration:
|
||||
|
||||
```yaml
|
||||
mcp_servers:
|
||||
playwright:
|
||||
command: npx
|
||||
args:
|
||||
- "@playwright/mcp@0.0.27"
|
||||
```
|
||||
|
||||
**Configuration Priority:** Command-line arguments > Configuration file > Environment variables > Default values
|
||||
|
||||
**Legacy JSON Configuration:** If using the older JSON format, see [docs/legacy_config.md](docs/legacy_config.md). We recommend migrating to YAML.
|
||||
|
||||
## 📖 Usage
|
||||
|
||||
### Basic Commands
|
||||
|
||||
```bash
|
||||
# Simple task execution
|
||||
trae-cli run "Create a hello world Python script"
|
||||
|
||||
# Check configuration
|
||||
trae-cli show-config
|
||||
|
||||
# Interactive mode
|
||||
trae-cli interactive
|
||||
```
|
||||
|
||||
### Provider-Specific Examples
|
||||
|
||||
```bash
|
||||
# OpenAI
|
||||
trae-cli run "Fix the bug in main.py" --provider openai --model gpt-4o
|
||||
|
||||
# Anthropic
|
||||
trae-cli run "Add unit tests" --provider anthropic --model claude-sonnet-4-20250514
|
||||
|
||||
# Google Gemini
|
||||
trae-cli run "Optimize this algorithm" --provider google --model gemini-2.5-flash
|
||||
|
||||
# OpenRouter (access to multiple providers)
|
||||
trae-cli run "Review this code" --provider openrouter --model "anthropic/claude-3-5-sonnet"
|
||||
trae-cli run "Generate documentation" --provider openrouter --model "openai/gpt-4o"
|
||||
|
||||
# Doubao
|
||||
trae-cli run "Refactor the database module" --provider doubao --model doubao-seed-1.6
|
||||
|
||||
# Ollama (local models)
|
||||
trae-cli run "Comment this code" --provider ollama --model qwen3
|
||||
```
|
||||
|
||||
### Advanced Options
|
||||
|
||||
```bash
|
||||
# Custom working directory
|
||||
trae-cli run "Add tests for utils module" --working-dir /path/to/project
|
||||
|
||||
# Save execution trajectory
|
||||
trae-cli run "Debug authentication" --trajectory-file debug_session.json
|
||||
|
||||
# Force patch generation
|
||||
trae-cli run "Update API endpoints" --must-patch
|
||||
|
||||
# Interactive mode with custom settings
|
||||
trae-cli interactive --provider openai --model gpt-4o --max-steps 30
|
||||
```
|
||||
|
||||
## Docker Mode Commands
|
||||
### Preparation
|
||||
**Important**: You need to make sure Docker is configured in your environment.
|
||||
|
||||
### Usage
|
||||
```bash
|
||||
# Specify a Docker image to run the task in a new container
|
||||
trae-cli run "Add tests for utils module" --docker-image python:3.11
|
||||
|
||||
# Specify a Docker image to run the task in a new container and mount the directory
|
||||
trae-cli run "write a script to print helloworld" --docker-image python:3.12 --working-dir test_workdir/
|
||||
|
||||
# Attach to an existing Docker container by ID (`--working-dir` is invalid with `--docker-container-id`)
|
||||
trae-cli run "Update API endpoints" --docker-container-id 91998a56056c
|
||||
|
||||
# Specify an absolute path to a Dockerfile to build an environment
|
||||
trae-cli run "Debug authentication" --dockerfile-path test_workspace/Dockerfile
|
||||
|
||||
# Specify a path to a local Docker image file (tar archive) to load
|
||||
trae-cli run "Fix the bug in main.py" --docker-image-file test_workspace/trae_agent_custom.tar
|
||||
|
||||
# Remove the Docker container after finishing the task (keep default)
|
||||
trae-cli run "Add tests for utils module" --docker-image python:3.11 --docker-keep false
|
||||
```
|
||||
|
||||
### Interactive Mode Commands
|
||||
|
||||
In interactive mode, you can use:
|
||||
- Type any task description to execute it
|
||||
- `status` - Show agent information
|
||||
- `help` - Show available commands
|
||||
- `clear` - Clear the screen
|
||||
- `exit` or `quit` - End the session
|
||||
|
||||
## 🛠️ Advanced Features
|
||||
|
||||
### Available Tools
|
||||
|
||||
Trae Agent provides a comprehensive toolkit for software engineering tasks including file editing, bash execution, structured thinking, and task completion. For detailed information about all available tools and their capabilities, see [docs/tools.md](docs/tools.md).
|
||||
|
||||
### Trajectory Recording
|
||||
|
||||
Trae Agent automatically records detailed execution trajectories for debugging and analysis:
|
||||
|
||||
```bash
|
||||
# Auto-generated trajectory file
|
||||
trae-cli run "Debug the authentication module"
|
||||
# Saves to: trajectories/trajectory_YYYYMMDD_HHMMSS.json
|
||||
|
||||
# Custom trajectory file
|
||||
trae-cli run "Optimize database queries" --trajectory-file optimization_debug.json
|
||||
```
|
||||
|
||||
Trajectory files contain LLM interactions, agent steps, tool usage, and execution metadata. For more details, see [docs/TRAJECTORY_RECORDING.md](docs/TRAJECTORY_RECORDING.md).
|
||||
|
||||
## 🔧 Development
|
||||
|
||||
### Contributing
|
||||
|
||||
For contribution guidelines, please refer to [CONTRIBUTING.md](CONTRIBUTING.md).
|
||||
|
||||
### Troubleshooting
|
||||
|
||||
**Import Errors:**
|
||||
```bash
|
||||
PYTHONPATH=. trae-cli run "your task"
|
||||
```
|
||||
|
||||
**API Key Issues:**
|
||||
```bash
|
||||
# Verify API keys
|
||||
echo $OPENAI_API_KEY
|
||||
trae-cli show-config
|
||||
```
|
||||
|
||||
**Command Not Found:**
|
||||
```bash
|
||||
uv run trae-cli run "your task"
|
||||
```
|
||||
|
||||
**Permission Errors:**
|
||||
```bash
|
||||
chmod +x /path/to/your/project
|
||||
```
|
||||
|
||||
## 📄 License
|
||||
|
||||
This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
|
||||
|
||||
## ✍️ Citation
|
||||
|
||||
```bibtex
|
||||
@article{traeresearchteam2025traeagent,
|
||||
title={Trae Agent: An LLM-based Agent for Software Engineering with Test-time Scaling},
|
||||
author={Trae Research Team and Pengfei Gao and Zhao Tian and Xiangxin Meng and Xinchen Wang and Ruida Hu and Yuanan Xiao and Yizhou Liu and Zhao Zhang and Junjie Chen and Cuiyun Gao and Yun Lin and Yingfei Xiong and Chao Peng and Xia Liu},
|
||||
year={2025},
|
||||
eprint={2507.23370},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.SE},
|
||||
url={https://arxiv.org/abs/2507.23370},
|
||||
}
|
||||
```
|
||||
|
||||
## 🙏 Acknowledgments
|
||||
|
||||
We thank Anthropic for building the [anthropic-quickstart](https://github.com/anthropics/anthropic-quickstarts) project that served as a valuable reference for the tool ecosystem.
|
||||
@@ -0,0 +1,7 @@
|
||||
# WeHub 来源说明
|
||||
|
||||
- 原始项目：`bytedance/trae-agent`
|
||||
- 原始仓库：https://github.com/bytedance/trae-agent
|
||||
- 导入方式：上游默认分支的最新快照
|
||||
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
|
||||
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
|
||||
@@ -0,0 +1,338 @@
|
||||
# Trajectory Recording Functionality
|
||||
|
||||
This document describes the trajectory recording functionality added to the Trae Agent project. The system captures detailed information about LLM interactions and agent execution steps for analysis, debugging, and auditing purposes.
|
||||
|
||||
## Overview
|
||||
|
||||
The trajectory recording system captures:
|
||||
|
||||
- **Raw LLM interactions**: Input messages, responses, token usage, and tool calls for various providers including Anthropic, OpenAI, Google Gemini, Azure, and others.
|
||||
- **Agent execution steps**: State transitions, tool calls, tool results, reflections, and errors
|
||||
- **Metadata**: Task description, timestamps, model configuration, and execution metrics
|
||||
|
||||
## Key Components
|
||||
|
||||
### 1. TrajectoryRecorder (`trae_agent/utils/trajectory_recorder.py`)
|
||||
|
||||
The core class that handles recording trajectory data to JSON files.
|
||||
|
||||
**Key methods:**
|
||||
|
||||
- `start_recording()`: Initialize recording with task metadata
|
||||
- `record_llm_interaction()`: Capture LLM request/response pairs
|
||||
- `record_agent_step()`: Capture agent execution steps
|
||||
- `finalize_recording()`: Complete recording and save final results
|
||||
|
||||
### 2. Client Integration
|
||||
|
||||
All supported LLM clients automatically record interactions when a trajectory recorder is attached.
|
||||
|
||||
**Anthropic Client** (`trae_agent/utils/anthropic_client.py`):
|
||||
|
||||
```python
|
||||
# Record trajectory if recorder is available
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_llm_interaction(
|
||||
messages=messages,
|
||||
response=llm_response,
|
||||
provider="anthropic",
|
||||
model=model_parameters.model,
|
||||
tools=tools
|
||||
)
|
||||
```
|
||||
|
||||
**OpenAI Client** (`trae_agent/utils/openai_client.py`):
|
||||
|
||||
```python
|
||||
# Record trajectory if recorder is available
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_llm_interaction(
|
||||
messages=messages,
|
||||
response=llm_response,
|
||||
provider="openai",
|
||||
model=model_parameters.model,
|
||||
tools=tools
|
||||
)
|
||||
```
|
||||
|
||||
**Google Gemini Client** (`trae_agent/utils/google_client.py`):
|
||||
|
||||
```python
|
||||
# Record trajectory if recorder is available
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_llm_interaction(
|
||||
messages=messages,
|
||||
response=llm_response,
|
||||
provider="google",
|
||||
model=model_parameters.model,
|
||||
tools=tools,
|
||||
)
|
||||
```
|
||||
|
||||
**Azure Client** (`trae_agent/utils/azure_client.py`):
|
||||
|
||||
```python
|
||||
# Record trajectory if recorder is available
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_llm_interaction(
|
||||
messages=messages,
|
||||
response=llm_response,
|
||||
provider="azure",
|
||||
model=model_parameters.model,
|
||||
tools=tools,
|
||||
)
|
||||
```
|
||||
|
||||
**Doubao Client** (`trae_agent/utils/doubao_client.py`):
|
||||
|
||||
```python
|
||||
# Record trajectory if recorder is available
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_llm_interaction(
|
||||
messages=messages,
|
||||
response=llm_response,
|
||||
provider="doubao",
|
||||
model=model_parameters.model,
|
||||
tools=tools,
|
||||
)
|
||||
```
|
||||
|
||||
**Ollama Client** (`trae_agent/utils/ollama_client.py`):
|
||||
|
||||
```python
|
||||
# Record trajectory if recorder is available
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_llm_interaction(
|
||||
messages=messages,
|
||||
response=llm_response,
|
||||
provider="openai", # Ollama client uses OpenAI's provider name for consistency
|
||||
model=model_parameters.model,
|
||||
tools=tools,
|
||||
)
|
||||
```
|
||||
|
||||
**OpenRouter Client** (`trae_agent/utils/openrouter_client.py`):
|
||||
|
||||
```python
|
||||
# Record trajectory if recorder is available
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_llm_interaction(
|
||||
messages=messages,
|
||||
response=llm_response,
|
||||
provider="openrouter",
|
||||
model=model_parameters.model,
|
||||
tools=tools,
|
||||
)
|
||||
```
|
||||
|
||||
### 3. Agent Integration
|
||||
|
||||
The base Agent class automatically records execution steps:
|
||||
|
||||
```python
|
||||
# Record agent step
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_agent_step(
|
||||
step_number=step.step_number,
|
||||
state=step.state.value,
|
||||
llm_messages=messages,
|
||||
llm_response=step.llm_response,
|
||||
tool_calls=step.tool_calls,
|
||||
tool_results=step.tool_results,
|
||||
reflection=step.reflection,
|
||||
error=step.error
|
||||
)
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### CLI Usage
|
||||
|
||||
#### Basic Recording (Auto-generated filename)
|
||||
|
||||
```bash
|
||||
trae run "Create a hello world Python script"
|
||||
# Trajectory saved to: trajectories/trajectory_20250612_220546.json
|
||||
```
|
||||
|
||||
#### Custom Filename
|
||||
|
||||
```bash
|
||||
trae run "Fix the bug in main.py" --trajectory-file my_debug_session.json
|
||||
# Trajectory saved to: my_debug_session.json
|
||||
```
|
||||
|
||||
#### Interactive Mode
|
||||
|
||||
```bash
|
||||
trae interactive --trajectory-file session.json
|
||||
```
|
||||
|
||||
### Programmatic Usage
|
||||
|
||||
```python
|
||||
from trae_agent.agent.trae_agent import TraeAgent
|
||||
from trae_agent.utils.llm_client import LLMProvider
|
||||
from trae_agent.utils.config import ModelParameters
|
||||
|
||||
# Create agent
|
||||
agent = TraeAgent(LLMProvider.ANTHROPIC, model_parameters, max_steps=10)
|
||||
|
||||
# Set up trajectory recording
|
||||
trajectory_path = agent.setup_trajectory_recording("my_trajectory.json")
|
||||
|
||||
# Configure and run task
|
||||
agent.new_task("My task", task_args)
|
||||
execution = await agent.execute_task()
|
||||
|
||||
# Trajectory is automatically saved
|
||||
print(f"Trajectory saved to: {trajectory_path}")
|
||||
```
|
||||
|
||||
## Trajectory File Format
|
||||
|
||||
The trajectory file is a JSON document with the following structure:
|
||||
|
||||
```json
|
||||
{
|
||||
"task": "Description of the task",
|
||||
"start_time": "2025-06-12T22:05:46.433797",
|
||||
"end_time": "2025-06-12T22:06:15.123456",
|
||||
"provider": "anthropic",
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"max_steps": 20,
|
||||
"llm_interactions": [
|
||||
{
|
||||
"timestamp": "2025-06-12T22:05:47.000000",
|
||||
"provider": "anthropic",
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"input_messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a software engineering assistant..."
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Create a hello world Python script"
|
||||
}
|
||||
],
|
||||
"response": {
|
||||
"content": "I'll help you create a hello world Python script...",
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"finish_reason": "end_turn",
|
||||
"usage": {
|
||||
"input_tokens": 150,
|
||||
"output_tokens": 75,
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"reasoning_tokens": null
|
||||
},
|
||||
"tool_calls": [
|
||||
{
|
||||
"call_id": "call_123",
|
||||
"name": "str_replace_based_edit_tool",
|
||||
"arguments": {
|
||||
"command": "create",
|
||||
"path": "hello.py",
|
||||
"file_text": "print('Hello, World!')"
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"tools_available": ["str_replace_based_edit_tool", "bash", "task_done"]
|
||||
}
|
||||
],
|
||||
"agent_steps": [
|
||||
{
|
||||
"step_number": 1,
|
||||
"timestamp": "2025-06-12T22:05:47.500000",
|
||||
"state": "thinking",
|
||||
"llm_messages": [...],
|
||||
"llm_response": {...},
|
||||
"tool_calls": [
|
||||
{
|
||||
"call_id": "call_123",
|
||||
"name": "str_replace_based_edit_tool",
|
||||
"arguments": {...}
|
||||
}
|
||||
],
|
||||
"tool_results": [
|
||||
{
|
||||
"call_id": "call_123",
|
||||
"success": true,
|
||||
"result": "File created successfully",
|
||||
"error": null
|
||||
}
|
||||
],
|
||||
"reflection": null,
|
||||
"error": null
|
||||
}
|
||||
],
|
||||
"success": true,
|
||||
"final_result": "Hello world Python script created successfully!",
|
||||
"execution_time": 28.689999
|
||||
}
|
||||
```
|
||||
|
||||
### Field Descriptions
|
||||
|
||||
**Root Level:**
|
||||
|
||||
- `task`: The original task description
|
||||
- `start_time`/`end_time`: ISO format timestamps
|
||||
- `provider`: LLM provider used (e.g., "anthropic", "openai", "google", "azure", "doubao", "ollama", "openrouter")
|
||||
- `model`: Model name
|
||||
- `max_steps`: Maximum allowed execution steps
|
||||
- `success`: Whether the task completed successfully
|
||||
- `final_result`: Final output or result message
|
||||
- `execution_time`: Total execution time in seconds
|
||||
|
||||
**LLM Interactions:**
|
||||
|
||||
- `timestamp`: When the interaction occurred
|
||||
- `provider`: LLM provider used for this interaction
|
||||
- `model`: Model used for this interaction
|
||||
- `input_messages`: Messages sent to the LLM
|
||||
- `response`: Complete LLM response including content, usage, and tool calls
|
||||
- `tools_available`: List of tools available during this interaction
|
||||
|
||||
**Agent Steps:**
|
||||
|
||||
- `step_number`: Sequential step number
|
||||
- `state`: Agent state ("thinking", "calling_tool", "reflecting", "completed", "error")
|
||||
- `llm_messages`: Messages used in this step
|
||||
- `llm_response`: LLM response for this step
|
||||
- `tool_calls`: Tools called in this step
|
||||
- `tool_results`: Results from tool execution
|
||||
- `reflection`: Agent's reflection on the step
|
||||
- `error`: Error message if the step failed
|
||||
|
||||
## Benefits
|
||||
|
||||
1. **Debugging**: Trace exactly what happened during agent execution
|
||||
2. **Analysis**: Understand LLM reasoning and tool usage patterns
|
||||
3. **Auditing**: Maintain records of what changes were made and why
|
||||
4. **Research**: Analyze agent behavior for improvements
|
||||
5. **Compliance**: Keep detailed logs of automated actions
|
||||
|
||||
## File Management
|
||||
|
||||
- Trajectory files are saved in the current working directory by default
|
||||
- Files use timestamp-based naming if no custom path is provided
|
||||
- Files are automatically created/overwritten
|
||||
- The system handles directory creation if needed
|
||||
- Files are saved continuously during execution (not just at the end)
|
||||
|
||||
## Security Considerations
|
||||
|
||||
- Trajectory files may contain sensitive information (API keys are not logged)
|
||||
- Store trajectory files securely if they contain proprietary code or data
|
||||
- Trajectory files are automatically saved to the `trajectories/` directory, which is excluded from version control
|
||||
|
||||
## Example Use Cases
|
||||
|
||||
1. **Debugging Failed Tasks**: Review what went wrong in agent execution
|
||||
2. **Performance Analysis**: Analyze token usage and execution patterns
|
||||
3. **Compliance Auditing**: Track all changes made by the agent
|
||||
4. **Model Comparison**: Compare behavior across different LLM providers/models
|
||||
5. **Tool Usage Analysis**: Understand which tools are used and how often
|
||||
@@ -0,0 +1,84 @@
|
||||
# Legacy JSON Configuration Guide
|
||||
|
||||
> **⚠️ DEPRECATED:** This JSON configuration format is deprecated and maintained for legacy compatibility only. For new installations, please use the [YAML configuration format](../README.md#configuration) instead.
|
||||
|
||||
## JSON Configuration Setup
|
||||
|
||||
**Configuration Setup:**
|
||||
|
||||
1. **Copy the example configuration file:**
|
||||
|
||||
```bash
|
||||
cp trae_config.json.example trae_config.json
|
||||
```
|
||||
|
||||
2. **Edit `trae_config.json` and replace the placeholder values with your actual credentials:**
|
||||
- Replace `"your_openai_api_key"` with your actual OpenAI API key
|
||||
- Replace `"your_anthropic_api_key"` with your actual Anthropic API key
|
||||
- Replace `"your_google_api_key"` with your actual Google API key
|
||||
- Replace `"your_azure_base_url"` with your actual Azure base URL
|
||||
- Replace other placeholder URLs and API keys as needed
|
||||
|
||||
**Note:** The `trae_config.json` file is ignored by git to prevent accidentally committing your API keys.
|
||||
|
||||
## JSON Configuration Structure
|
||||
|
||||
Trae Agent uses a JSON configuration file for settings. Please refer to the `trae_config.json.example` file in the root directory for the detailed configuration structure.
|
||||
|
||||
**Configuration Priority:**
|
||||
|
||||
1. Command-line arguments (highest)
|
||||
2. Configuration file values
|
||||
3. Environment variables
|
||||
4. Default values (lowest)
|
||||
|
||||
## Example JSON Configuration
|
||||
|
||||
The JSON configuration file contains provider-specific settings for various LLM services:
|
||||
|
||||
```json
|
||||
{
|
||||
"default_provider": "anthropic",
|
||||
"max_steps": 20,
|
||||
"enable_lakeview": true,
|
||||
"model_providers": {
|
||||
"openai": {
|
||||
"api_key": "your_openai_api_key",
|
||||
"base_url": "https://api.openai.com/v1",
|
||||
"model": "gpt-4o",
|
||||
"max_tokens": 128000,
|
||||
"temperature": 0.5,
|
||||
"top_p": 1,
|
||||
"max_retries": 10
|
||||
},
|
||||
"anthropic": {
|
||||
"api_key": "your_anthropic_api_key",
|
||||
"base_url": "https://api.anthropic.com",
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"max_tokens": 4096,
|
||||
"temperature": 0.5,
|
||||
"top_p": 1,
|
||||
"top_k": 0,
|
||||
"max_retries": 10
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Migration to YAML
|
||||
|
||||
To migrate from JSON to YAML configuration:
|
||||
|
||||
1. **Create a new YAML configuration file:**
|
||||
```bash
|
||||
cp trae_config.yaml.example trae_config.yaml
|
||||
```
|
||||
|
||||
2. **Transfer your settings** from `trae_config.json` to `trae_config.yaml` following the new structure
|
||||
|
||||
3. **Remove the old JSON file** (optional but recommended):
|
||||
```bash
|
||||
rm trae_config.json
|
||||
```
|
||||
|
||||
For detailed YAML configuration instructions, please refer to the main [README.md](../README.md#configuration).
|
||||
@@ -0,0 +1,96 @@
|
||||
# Trae Agent Roadmap
|
||||
|
||||
This roadmap outlines the planned features and enhancements for Trae Agent. Our goal is to build a comprehensive, research-friendly AI agent platform that serves both developers and researchers in the rapidly evolving field of AI agents.
|
||||
|
||||
## SDK Development
|
||||
|
||||
### Overview
|
||||
Develop a comprehensive Software Development Kit (SDK) to enable programmatic access to Trae Agent capabilities, making it easier for developers to integrate agent functionality into their applications and workflows.
|
||||
|
||||
### Key Features
|
||||
- **Headless Interface**: Programmatic API for agent interaction without CLI dependency
|
||||
- **Streamed Trajectory Recording**: Real-time access to detailed LLM interactions and tool execution data
|
||||
|
||||
### Benefits
|
||||
- **Developer Integration**: Enables seamless integration of Trae Agent into existing applications, CI/CD pipelines, and development workflows
|
||||
- **Real-time Monitoring**: Streamed trajectory recording allows for live monitoring of agent behavior, enabling immediate feedback and intervention when needed
|
||||
- **Automation**: Facilitates automated testing, batch processing, and unattended agent operations
|
||||
- **Research Applications**: Provides researchers with programmatic access to agent internals for studying agent behavior and conducting experiments
|
||||
|
||||
## Sandbox Environment
|
||||
|
||||
### Overview
|
||||
Implement secure sandbox environments for task execution, providing isolated and controlled environments where agents can operate safely without affecting the host system.
|
||||
|
||||
### Key Features
|
||||
- **Isolated Task Execution**: Run agent tasks within containerized or virtualized environments
|
||||
- **Parallel Task Execution**: Support for running multiple agent instances simultaneously
|
||||
|
||||
### Benefits
|
||||
- **Security**: Protects the host system from potentially harmful operations during agent execution
|
||||
- **Reproducibility**: Ensures consistent execution environments across different systems and deployments
|
||||
- **Scalability**: Parallel execution capabilities enable handling multiple tasks simultaneously, improving throughput
|
||||
- **Development Safety**: Allows safe experimentation with agent behavior without risk to production systems
|
||||
- **Multi-tenancy**: Enables serving multiple users or projects with isolated agent instances
|
||||
|
||||
## Trajectory Analysis
|
||||
|
||||
### Overview
|
||||
Enhance trajectory recording and analysis capabilities by integrating with popular machine learning operations (MLOps) platforms and providing advanced analytics tools.
|
||||
|
||||
### Key Features
|
||||
- **MLOps Integration**: Connect with backends such as Weights & Biases (Wandb) Weave and MLFlow
|
||||
- **Advanced Analytics**: Provide detailed insights into agent performance, token usage, and decision patterns
|
||||
|
||||
### Benefits
|
||||
- **Performance Optimization**: Detailed analytics help identify bottlenecks and optimization opportunities in agent workflows
|
||||
- **Research Insights**: Rich trajectory data enables researchers to study agent behavior patterns, decision-making processes, and tool usage
|
||||
- **Debugging & Troubleshooting**: Enhanced logging and visualization make it easier to diagnose issues and understand agent failures
|
||||
- **Model Comparison**: Integration with MLOps platforms allows for systematic comparison of different models and configurations
|
||||
- **Compliance & Auditing**: Comprehensive logging supports audit requirements and regulatory compliance needs
|
||||
|
||||
## Tools and Model Context Protocol (MCP)
|
||||
|
||||
### Overview
|
||||
Expand the tool ecosystem to support more file formats and integrate with the Model Context Protocol (MCP) for enhanced interoperability and standardized tool interfaces.
|
||||
|
||||
### Key Features
|
||||
- **Structured File Support**: Enhanced support for Jupyter Notebooks, configuration files, and other structured formats
|
||||
- **MCP Integration**: Implement Model Context Protocol for standardized tool communication
|
||||
|
||||
### Benefits
|
||||
- **Enhanced Productivity**: Better support for Jupyter Notebooks enables seamless data science and research workflows
|
||||
- **Standardization**: MCP adoption ensures compatibility with other AI tools and platforms
|
||||
- **Extensibility**: Standardized interfaces make it easier for third-party developers to create and share tools
|
||||
- **Ecosystem Growth**: MCP support opens access to a broader ecosystem of existing tools and services
|
||||
- **Interoperability**: Seamless integration with other MCP-compatible AI systems and workflows
|
||||
|
||||
## Advanced Agentic Flows and Multi-Agent Support
|
||||
|
||||
### Overview
|
||||
Develop sophisticated agent orchestration capabilities, including support for multiple specialized agents working together and advanced workflow patterns.
|
||||
|
||||
### Key Features
|
||||
- **Multi-Agent Coordination**: Support for multiple agents collaborating on complex tasks
|
||||
- **Advanced Workflow Patterns**: Implement sophisticated agentic flows beyond simple linear task execution
|
||||
- **Agent Specialization**: Enable creation of specialized agents for specific domains or tasks
|
||||
|
||||
### Benefits
|
||||
- **Complex Problem Solving**: Multi-agent systems can tackle problems that require diverse expertise and parallel processing
|
||||
- **Scalability**: Distributed agent architecture enables handling larger and more complex projects
|
||||
- **Specialization**: Domain-specific agents can provide deeper expertise in particular areas (e.g., frontend development, data analysis, security)
|
||||
- **Robustness**: Multi-agent systems can provide redundancy and fault tolerance
|
||||
- **Research Opportunities**: Advanced agentic flows enable research into agent communication, coordination, and emergent behaviors
|
||||
|
||||
## Community Involvement
|
||||
|
||||
We encourage community participation in shaping this roadmap. Please:
|
||||
|
||||
- **Submit feature requests**: Share your ideas and use cases through GitHub issues
|
||||
- **Contribute to discussions**: Participate in roadmap discussions and RFC processes
|
||||
- **Contribute code**: Help implement features that align with your needs and expertise
|
||||
- **Share research**: Contribute findings and insights from your research with Trae Agent
|
||||
|
||||
---
|
||||
|
||||
*This roadmap is a living document that will evolve based on community needs, research developments, and technological advances in the AI agent space.*
|
||||
@@ -0,0 +1,85 @@
|
||||
# Tools
|
||||
|
||||
Trae Agent provides five built-in tools for software engineering tasks:
|
||||
|
||||
## str_replace_based_edit_tool
|
||||
|
||||
File and directory manipulation tool with persistent state.
|
||||
|
||||
**Operations:**
|
||||
- `view` - Display file contents with line numbers, or list directory contents up to 2 levels deep
|
||||
- `create` - Create new files (fails if file already exists)
|
||||
- `str_replace` - Replace exact string matches in files (must be unique)
|
||||
- `insert` - Insert text after a specified line number
|
||||
|
||||
**Key features:**
|
||||
- Requires absolute paths (e.g., `/repo/file.py`)
|
||||
- String replacements must match exactly, including whitespace
|
||||
- Supports line range viewing for large files
|
||||
|
||||
## bash
|
||||
|
||||
Execute shell commands in a persistent session.
|
||||
|
||||
**Features:**
|
||||
- Commands run in a shared bash session that maintains state
|
||||
- 120-second timeout per command
|
||||
- Session restart capability
|
||||
- Background process support
|
||||
|
||||
**Usage notes:**
|
||||
- Use `restart: true` to reset the session
|
||||
- Avoid commands with excessive output
|
||||
- Long-running commands should use `&` for background execution
|
||||
|
||||
## sequential_thinking
|
||||
|
||||
Structured problem-solving tool for complex analysis.
|
||||
|
||||
**Capabilities:**
|
||||
- Break down problems into sequential thoughts
|
||||
- Revise and branch from previous thoughts
|
||||
- Dynamically adjust the number of thoughts needed
|
||||
- Track thinking history and alternative approaches
|
||||
- Generate and verify solution hypotheses
|
||||
|
||||
**Parameters:**
|
||||
- `thought` - Current thinking step
|
||||
- `thought_number` / `total_thoughts` - Progress tracking
|
||||
- `next_thought_needed` - Continue thinking flag
|
||||
- `is_revision` / `revises_thought` - Revision tracking
|
||||
- `branch_from_thought` / `branch_id` - Alternative exploration
|
||||
|
||||
## task_done
|
||||
|
||||
Signal task completion with verification requirement.
|
||||
|
||||
**Purpose:**
|
||||
- Mark tasks as successfully completed
|
||||
- Must be called only after proper verification
|
||||
- Encourages writing test/reproduction scripts
|
||||
|
||||
**Output:**
|
||||
- Simple "Task done." message
|
||||
- No parameters required
|
||||
|
||||
## json_edit_tool
|
||||
|
||||
Precise JSON file editing using JSONPath expressions.
|
||||
|
||||
**Operations:**
|
||||
- `view` - Display entire file or content at specific JSONPaths
|
||||
- `set` - Update existing values at specified paths
|
||||
- `add` - Add new properties to objects or append to arrays
|
||||
- `remove` - Delete elements at specified paths
|
||||
|
||||
**JSONPath examples:**
|
||||
- `$.users[0].name` - First user's name
|
||||
- `$.config.database.host` - Nested object property
|
||||
- `$.items[*].price` - All item prices
|
||||
- `$..key` - Recursive search for key
|
||||
|
||||
**Features:**
|
||||
- Validates JSON syntax and structure
|
||||
- Preserves formatting with pretty printing option
|
||||
- Detailed error messages for invalid operations
|
||||
@@ -0,0 +1,236 @@
|
||||
# Evaluation for Trae Agent
|
||||
|
||||
This document explains how to evaluate [Trae Agent](https://github.com/bytedance/trae-agent) using [SWE-bench](https://www.swebench.com/), [SWE-bench-Live](https://swe-bench-live.github.io/), and [Multi-SWE-bench](https://multi-swe-bench.github.io/).
|
||||
|
||||
## Overview
|
||||
|
||||
**SWE-bench** is a benchmark that evaluates language models on real-world software engineering tasks. It contains GitHub issues from popular Python repositories that have been solved by human developers. The benchmark evaluates whether an agent can generate the correct patch to fix the issue.
|
||||
|
||||
**SWE-bench-Live** is a live benchmark for issue resolving, designed to evaluate an AI system's ability to complete real-world software engineering tasks. Thanks to our automated dataset curation pipeline, we plan to update SWE-bench-Live on a monthly basis to provide the community with up-to-date task instances and support rigorous and contamination-free evaluation.
|
||||
|
||||
**Multi-SWE-bench** is a multilingual benchmark for issue resolving. It spans 7 languages (i.e., Java, TypeScript, JavaScript, Go, Rust, C, and C++) with 1,632 high-quality instances, curated from 2,456 candidates by 68 expert annotators for reliability.
|
||||
|
||||
The evaluation process involves:
|
||||
1. **Setup**: Preparing the evaluation environment with Docker containers
|
||||
2. **Execution**: Running Trae Agent on instances to generate patches
|
||||
3. **Evaluation**: Testing the generated patches against the ground truth using harness
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Before running the evaluation, ensure you have:
|
||||
|
||||
- **Docker**: Required for containerized evaluation environments
|
||||
- **Python 3.12+**: For running the evaluation scripts
|
||||
- **Git**: For cloning repositories
|
||||
- **Sufficient disk space**: Docker images can be several GBs per instance
|
||||
- **API Keys**: OpenAI/Anthropic API keys for Trae Agent
|
||||
|
||||
## Setup Instructions
|
||||
|
||||
Make sure installing extra dependencies for evaluation and running scripts in the `evaluation` directory.
|
||||
|
||||
```bash
|
||||
uv sync --extra evaluation
|
||||
cd evaluation
|
||||
```
|
||||
|
||||
### 1. Clone and Setup Benchmark Harness
|
||||
|
||||
The `setup.sh` script automates the setup of benchmark harness:
|
||||
|
||||
```bash
|
||||
chmod +x setup.sh
|
||||
./setup.sh [swe_bench|swe_bench_live|multi_swe_bench]
|
||||
```
|
||||
|
||||
- `swe_bench`: Setup for SWE-Bench
|
||||
- `swe_bench_live`: Setup for SWE-Bench-Live
|
||||
- `multi_swe_bench`: Setup for Multi-SWE-Bench
|
||||
|
||||
This script:
|
||||
- Clones the benchmark repository
|
||||
- Checks out a specific commit for reproducibility (it is the most recent commit hash at the time of writing this document.)
|
||||
- Creates a Python virtual environment
|
||||
- Installs the benchmark harness
|
||||
|
||||
### 2. Configure Trae Agent
|
||||
|
||||
Ensure your `trae_config.yaml` file is properly configured with valid API keys:
|
||||
|
||||
```
|
||||
agents:
|
||||
trae_agent:
|
||||
enable_lakeview: false
|
||||
model: trae_agent_model # the model configuration name for Trae Agent
|
||||
max_steps: 200 # max number of agent steps
|
||||
tools: # tools used with Trae Agent
|
||||
- bash
|
||||
- str_replace_based_edit_tool
|
||||
- sequentialthinking
|
||||
- task_done
|
||||
|
||||
model_providers: # model providers configuration
|
||||
anthropic:
|
||||
api_key: your_anthropic_api_key
|
||||
provider: anthropic
|
||||
openai:
|
||||
api_key: your_openai_api_key
|
||||
provider: openai
|
||||
|
||||
models:
|
||||
trae_agent_model:
|
||||
model_provider: anthropic
|
||||
model: claude-sonnet-4-20250514
|
||||
max_tokens: 4096
|
||||
temperature: 0.5
|
||||
top_p: 0.9
|
||||
top_k: 40
|
||||
max_retries: 1
|
||||
parallel_tool_calls: 1
|
||||
```
|
||||
|
||||
### 3. Optional: Docker Environment Configuration
|
||||
|
||||
Create a `docker_env_config.json` file if you need custom environment variables:
|
||||
|
||||
```json
|
||||
{
|
||||
"preparation_env": {
|
||||
"HTTP_PROXY": "http://proxy.example.com:8080",
|
||||
"HTTPS_PROXY": "https://proxy.example.com:8080"
|
||||
},
|
||||
"experiment_env": {
|
||||
"CUSTOM_VAR": "value"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## Usage
|
||||
|
||||
### Basic Usage
|
||||
The evaluation script `run_evaluation.py` provides several modes of operation:
|
||||
|
||||
```bash
|
||||
# Run evaluation on all instances of SWE-bench_Verified
|
||||
python run_evaluation.py --dataset SWE-bench_Verified --working-dir ./trae-workspace
|
||||
|
||||
# Run evaluation on specific instances
|
||||
python run_evaluation.py --instance_ids django__django-12345 scikit-learn__scikit-learn-67890
|
||||
|
||||
# Run with custom configuration
|
||||
python run_evaluation.py --config-file trae_config.yaml --run-id experiment-1
|
||||
```
|
||||
|
||||
### Available Benchmarks and Datasets
|
||||
|
||||
**SWE-bench**
|
||||
- **SWE-bench_Verified**
|
||||
- **SWE-bench_Lite**
|
||||
- **SWE-bench**
|
||||
|
||||
**SWE-bench-Live**:
|
||||
- **SWE-bench-Live/lite**
|
||||
- **SWE-bench-Live/verified**
|
||||
- **SWE-bench-Live/full**
|
||||
|
||||
**Multi-SWE-bench**:
|
||||
- **Multi-SWE-bench-flash** (Please download `multi_swe_bench_flash.jsonl` from https://huggingface.co/datasets/ByteDance-Seed/Multi-SWE-bench-flash/tree/main and place it in the `evaluation` directory.)
|
||||
- **Multi-SWE-bench_mini** (Please download `multi_swe_bench_mini.jsonl` from https://huggingface.co/datasets/ByteDance-Seed/Multi-SWE-bench_mini/tree/main and place it in the `evaluation` directory.)
|
||||
|
||||
### Evaluation Modes
|
||||
|
||||
The script supports three modes:
|
||||
|
||||
1. **`expr`** (Expression only): Generate patches without evaluation
|
||||
2. **`eval`** (Evaluation only): Evaluate existing patches
|
||||
3. **`e2e`** (End-to-end): Both generate and evaluate patches (default)
|
||||
|
||||
```bash
|
||||
# Only generate patches
|
||||
python run_evaluation.py --mode expr --dataset SWE-bench_Verified
|
||||
|
||||
# Only evaluate existing patches
|
||||
python run_evaluation.py --mode eval --benchmark-harness-path ./SWE-bench
|
||||
|
||||
# End-to-end evaluation (default)
|
||||
python swebench.py --mode e2e --benchmark-harness-path ./SWE-bench
|
||||
```
|
||||
|
||||
### Full Command Reference
|
||||
|
||||
```bash
|
||||
python run_evaluation.py \
|
||||
--benchmark SWE-bench \
|
||||
--dataset SWE-bench_Verified \
|
||||
--config-file ./trae_config.yaml \
|
||||
--run-id experiment-1 \
|
||||
--benchmark-harness-path ./SWE-bench \
|
||||
--docker-env-config ./docker_env_config.json \
|
||||
--mode e2e \
|
||||
--max_workers 4 \
|
||||
--instance_ids astropy__astropy-13453
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
- `--benchmark`: Benchmark to use
|
||||
- `--dataset`: Dataset to use
|
||||
- `--config-file`: Trae Agent configuration file
|
||||
- `--run-id`: Run ID for benchmark evaluation
|
||||
- `--benchmark-harness-path`: Path to SWE-bench harness (required for evaluation)
|
||||
- `--docker-env-config`: Docker environment configuration file
|
||||
- `--mode`: Evaluation mode (`e2e`, `expr`, `eval`)
|
||||
- `--max_workers`: Maximum number of worker processes to use for parallel execution.
|
||||
- `--instance_ids`: Instances to use
|
||||
|
||||
## How It Works
|
||||
|
||||
### 1. Image Preparation
|
||||
|
||||
The script first checks for required Docker images:
|
||||
- Each instance has a specific Docker image
|
||||
- Images are pulled automatically if not present locally
|
||||
- Base Ubuntu image is used for preparing Trae Agent
|
||||
|
||||
### 2. Trae Agent Preparation
|
||||
|
||||
The script builds Trae Agent in a Docker container:
|
||||
- Creates artifacts (`trae-agent.tar`, `uv.tar`, `uv_shared.tar`)
|
||||
- These artifacts are reused across all instances for efficiency
|
||||
|
||||
### 3. Instance Execution
|
||||
|
||||
For each instance:
|
||||
1. **Container Setup**: Prepares a Docker container with the instance's environment
|
||||
2. **Problem Statement**: Writes the GitHub issue description to a file
|
||||
3. **Trae Agent Execution**: Runs Trae Agent to generate a patch
|
||||
4. **Patch Collection**: Saves the generated patch for evaluation
|
||||
|
||||
### 4. Evaluation
|
||||
|
||||
Using benchmark harness:
|
||||
1. **Patch Collection**: Collects all generated patches into `predictions.json`
|
||||
2. **Test Execution**: Runs the patches against test suites in Docker containers
|
||||
3. **Result Generation**: Produces evaluation results with pass/fail status
|
||||
|
||||
## Understanding Results
|
||||
|
||||
### Output Files
|
||||
|
||||
The evaluation creates several files in the working directory:
|
||||
|
||||
```
|
||||
results/{benchmark}_{dataset}_{run_id}/
|
||||
├── predictions.json # Generated patches for evaluation
|
||||
├── results.json # Final evaluation results
|
||||
├── {instance_id}/ # Folder for each instance
|
||||
│ ├── problem_statement.txt # GitHub issue description
|
||||
│ ├── {instance_id}.patch # Generated patch
|
||||
│ ├── {instance_id}.json # Trajectory file
|
||||
│ └── ...
|
||||
trae-workspace/
|
||||
├── trae_config.yaml # Trae Agent configuration file
|
||||
├── trae-agent.tar # Trae Agent build artifacts
|
||||
├── uv.tar # UV binary
|
||||
└── uv_shared.tar # UV shared files
|
||||
```
|
||||
@@ -0,0 +1,2 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
@@ -0,0 +1,102 @@
|
||||
# Selector Agent
|
||||
|
||||
This document explains how to further enhance [Trae Agent](https://github.com/bytedance/trae-agent) using the selector agent.
|
||||
Selector agent is the first agent-based ensemble reasoning approach for repository-level issue resolution.
|
||||
It formulates our goal as an optimal solution search problem and addresses two key challenges, i.e., large ensemble spaces and repository-level understanding, through modular agents for generation, pruning, and selection.
|
||||
|
||||
## 📖 Demo
|
||||
|
||||
### Regression Testing
|
||||
For regression testing, please refer to [Agentless](https://github.com/OpenAutoCoder/Agentless/blob/main/README_swebench.md).
|
||||
|
||||
Each result entry contains a `regression` field that indicates test outcomes:
|
||||
- An empty array [] signifies the patch successfully passed all regression tests;
|
||||
- Any non-empty value indicates the patch caused test failures (with details specifying which tests failed).
|
||||
|
||||
### Preparation
|
||||
|
||||
**Important:** You need to download a Python 3.12 package from [Google Drive](https://drive.google.com/file/d/1dF7kbcmdLRJu7TEh8G7Oe8_6NY3aieKa/view?usp=sharing) and unzip it into `evaluation/patch_selection/trae_selector/tools/py312`. This is used to execute agent tools in Docker containers.
|
||||
|
||||
### Input Format
|
||||
|
||||
Patch candidates are stored in a JSON line file. For each instance, the structure is as follows:
|
||||
|
||||
```json
|
||||
{
|
||||
"instance_id": "django__django-14017",
|
||||
"issue": "Issue description....",
|
||||
"patches": [
|
||||
"patch diff 1",
|
||||
"patch diff 2",
|
||||
...,
|
||||
"patch diff N",
|
||||
],
|
||||
"success_id": [
|
||||
1,
|
||||
0,
|
||||
...,
|
||||
1
|
||||
],
|
||||
"regressions": [
|
||||
[regression_test_names for patch diff 1..],
|
||||
[regression_test_names for patch diff 2..],
|
||||
...,
|
||||
[regression_test_names for patch diff N..],
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
Note: success_id is either 1 (the corresponding patch diff is a correct patch) or 0 (the corresponding patch diff is a wrong patch). Once a patch is selected by the Selector Agent, we can quickly report if the selected patch is correct or not.
|
||||
|
||||
The regressions field is optional. If you have done regression test selection using Agentless, you can fill in selected regression tests here.
|
||||
|
||||
### Patch Selection
|
||||
|
||||
```bash
|
||||
python3 evaluation/patch_selection/selector.py \
|
||||
--instances_path "path/to/swebench-verified.json" \
|
||||
--candidate_path "path/to/patch_candidates.jsonl" \
|
||||
--result_path "path/to/save/results" \
|
||||
--num_candidate NUMBER_OF_PATCH_CANDIDATES_PER_INSTANCE \
|
||||
--max_workers 10 \
|
||||
--group_size GROUP_SIZE \
|
||||
--max_retry 20 \
|
||||
--max_turn 200 \
|
||||
--config_file trae_config.yaml \
|
||||
--model_name MODEL_NAME_IN_CONFIG_FILE \
|
||||
--majority_voting
|
||||
```
|
||||
|
||||
Note: if you have a lot of patch candidates, for example 50, you can set group_size to 10. The patch selection is done by 5 (50/10) groups. A patch is selected for each group. You can then select from these 5.
|
||||
|
||||
`--majority_voting` is optional. If enabled, for each candidate group, multiple patch selection is conducted and the patch with most selected frequency is the final answer. This mode consumes more token consumption.
|
||||
|
||||
### Example
|
||||
|
||||
After running with [example.jsonl](example/example.jsonl), in the result_path, we get the following files:
|
||||
|
||||
```text
|
||||
├── log
|
||||
│ └── group_0
|
||||
│ └── astropy__astropy-14369_voting_0_trail_1.json
|
||||
├── output
|
||||
│ └── group_0
|
||||
│ └── astropy__astropy-14369.log
|
||||
├── patch
|
||||
│ └── group_0
|
||||
│ └── astropy__astropy-14369_1.patch
|
||||
└── statistics
|
||||
└── group_0
|
||||
└── astropy__astropy-14369.json
|
||||
```
|
||||
|
||||
* The file in the log directory stores LLM interaction history.
|
||||
* The file in the output directory stores raw standard output and standard error.
|
||||
* Patch directory stores selected patches.
|
||||
* Statistics directory stores whether the selected patch is correct or not.
|
||||
|
||||
You can use the `analysis.py` script to visualise the selection results (even during the selection is running to see intermediate results)
|
||||
|
||||
```bash
|
||||
python3 analysis.py --output_path "path/to/save/results"
|
||||
```
|
||||
@@ -0,0 +1,155 @@
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import os
|
||||
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--output_path", type=str, required=True)
|
||||
parser.add_argument("--group_id", type=int, required=False, default=None)
|
||||
args = parser.parse_args()
|
||||
|
||||
output_path = args.output_path
|
||||
statistics_path = output_path + "/statistics"
|
||||
|
||||
if args.group_id is not None:
|
||||
statistics_folder_path = statistics_path + f"/group_{args.group_id}"
|
||||
result = {f"group_{args.group_id}": analyze_group(statistics_folder_path)}
|
||||
else:
|
||||
# get all groups in the statistics directory
|
||||
group_ids = [
|
||||
f
|
||||
for f in os.listdir(statistics_path)
|
||||
if os.path.isdir(os.path.join(statistics_path, f))
|
||||
]
|
||||
result = {}
|
||||
for group_id in group_ids:
|
||||
statistics_folder_path = statistics_path + f"/{group_id}"
|
||||
result[f"{group_id}"] = analyze_group(statistics_folder_path)
|
||||
|
||||
# sort result by success_rate_among_all
|
||||
result = dict(
|
||||
sorted(result.items(), key=lambda item: item[1]["success_rate_among_all"], reverse=True)
|
||||
)
|
||||
|
||||
table = Table(title=f"Statistics for Selector Experiment {output_path}")
|
||||
# save to csv
|
||||
with open(output_path + "/analysis.csv", "w") as f:
|
||||
writer = csv.writer(f)
|
||||
table_header = [
|
||||
"group_id",
|
||||
"total",
|
||||
"completion_rate",
|
||||
"all_success",
|
||||
"all_failed",
|
||||
"need_to_select",
|
||||
"success_selection",
|
||||
"success_selection_in_need_to_select",
|
||||
"success_rate_in_need_to_select",
|
||||
"success_rate_among_all",
|
||||
]
|
||||
for header in table_header:
|
||||
if header == "success_rate_in_need_to_select":
|
||||
table.add_column(header, justify="right", no_wrap=True, style="cyan")
|
||||
elif header == "success_rate_among_all":
|
||||
table.add_column(header, justify="right", no_wrap=True, style="magenta")
|
||||
else:
|
||||
table.add_column(header, justify="right", no_wrap=True)
|
||||
writer.writerow(table_header)
|
||||
|
||||
max_success_rate_in_need_to_select = 0
|
||||
max_success_rate_group_id = ""
|
||||
max_success_rate_among_all = 0
|
||||
max_success_rate_among_all_group_id = ""
|
||||
table_rows = []
|
||||
for group_id, record in result.items():
|
||||
row = [
|
||||
group_id,
|
||||
record["total"],
|
||||
record["completion_rate"],
|
||||
record["all_success"],
|
||||
record["all_failed"],
|
||||
record["need_to_select"],
|
||||
record["success_selection"],
|
||||
record["success_selection_in_need_to_select"],
|
||||
record["success_rate_in_need_to_select"],
|
||||
record["success_rate_among_all"],
|
||||
]
|
||||
|
||||
# make the largest success rate in need to select and success rate among all bold
|
||||
if float(record["success_rate_in_need_to_select"]) > max_success_rate_in_need_to_select:
|
||||
max_success_rate_in_need_to_select = float(record["success_rate_in_need_to_select"])
|
||||
max_success_rate_group_id = group_id
|
||||
if float(record["success_rate_among_all"]) > max_success_rate_among_all:
|
||||
max_success_rate_among_all = float(record["success_rate_among_all"])
|
||||
max_success_rate_among_all_group_id = group_id
|
||||
table_rows.append(row)
|
||||
writer.writerow(row)
|
||||
|
||||
for row in table_rows:
|
||||
if row[0] == max_success_rate_group_id:
|
||||
row[8] = f"[strong][underline]{row[8] * 100:.2f}%[/underline][/strong]"
|
||||
if row[0] == max_success_rate_among_all_group_id:
|
||||
row[9] = f"[strong][underline]{row[9] * 100:.2f}%[/underline][/strong]"
|
||||
for i in range(len(row)):
|
||||
if isinstance(row[i], float):
|
||||
row[i] = f"{row[i] * 100:.2f}%"
|
||||
else:
|
||||
row[i] = str(row[i])
|
||||
table.add_row(*row)
|
||||
|
||||
# print in table
|
||||
console = Console()
|
||||
console.print(table)
|
||||
|
||||
|
||||
def analyze_group(statistics_folder_path, total_num_instances=500):
|
||||
all_success = 0
|
||||
all_failed = 0
|
||||
need_to_select = 0
|
||||
success_selection = 0
|
||||
success_selection_in_need_to_select = 0
|
||||
total = 0
|
||||
|
||||
# list all json files in the statistics folder
|
||||
json_files = [f for f in os.listdir(statistics_folder_path) if f.endswith(".json")]
|
||||
for json_file in json_files:
|
||||
with open(os.path.join(statistics_folder_path, json_file), "r") as f:
|
||||
try:
|
||||
data = json.loads(f.read())
|
||||
except Exception:
|
||||
print(f"Error loading {os.path.join(statistics_folder_path, json_file)}")
|
||||
if data["is_all_success"]:
|
||||
all_success += 1
|
||||
if data["is_all_failed"]:
|
||||
all_failed += 1
|
||||
if not data["is_all_success"] and not data["is_all_failed"]:
|
||||
need_to_select += 1
|
||||
if data["is_success"] == 1:
|
||||
success_selection_in_need_to_select += 1
|
||||
if data["is_success"] == 1:
|
||||
success_selection += 1
|
||||
total += 1
|
||||
|
||||
return {
|
||||
"total": total,
|
||||
"completion_rate": float(total) / float(total_num_instances),
|
||||
"all_success": all_success,
|
||||
"all_failed": all_failed,
|
||||
"need_to_select": need_to_select,
|
||||
"success_selection": success_selection,
|
||||
"success_selection_in_need_to_select": success_selection_in_need_to_select,
|
||||
"success_rate_in_need_to_select": float(success_selection_in_need_to_select)
|
||||
/ float(need_to_select)
|
||||
if need_to_select > 0
|
||||
else 0,
|
||||
"success_rate_among_all": float(success_selection) / float(total) if total > 0 else 0,
|
||||
}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,102 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from trae_selector.selector_evaluation import SelectorEvaluation
|
||||
|
||||
from trae_agent.utils.config import Config
|
||||
|
||||
_ = load_dotenv() # take environment variables
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
_ = parser.add_argument(
|
||||
"--instances_path",
|
||||
default="swe_bench/swebench-verified.json",
|
||||
help="Path to instances JSON file",
|
||||
)
|
||||
_ = parser.add_argument("--candidate_path", required=True, help="Path to candidate patches")
|
||||
_ = parser.add_argument("--result_path", required=True, help="Path to save results")
|
||||
_ = parser.add_argument(
|
||||
"--num_candidate", type=int, default=10, help="The number of candidate patches"
|
||||
)
|
||||
_ = parser.add_argument("--max_workers", type=int, default=10, help="Max number of workers")
|
||||
_ = parser.add_argument(
|
||||
"--group_size", type=int, default=10, help="Group size of candidate patches"
|
||||
)
|
||||
_ = parser.add_argument(
|
||||
"--max_retry", type=int, default=3, help="Max retry times of LLM responses"
|
||||
)
|
||||
_ = parser.add_argument(
|
||||
"--max_turn", type=int, default=50, help="Max turn times of Selector Agent"
|
||||
)
|
||||
_ = parser.add_argument("--majority_voting", action=argparse.BooleanOptionalAction)
|
||||
_ = parser.add_argument(
|
||||
"--config_file", type=str, default="config.yaml", help="Path to config file"
|
||||
)
|
||||
_ = parser.add_argument("--model_name", type=str, default="default_model", help="Model name")
|
||||
args = parser.parse_args()
|
||||
args.log_path = os.path.join(args.result_path, "log")
|
||||
args.output_path = os.path.join(args.result_path, "output")
|
||||
args.patches_path = os.path.join(args.result_path, "patch")
|
||||
args.statistics_path = os.path.join(args.result_path, "statistics")
|
||||
[
|
||||
os.makedirs(_)
|
||||
for _ in [args.log_path, args.patches_path, args.output_path, args.statistics_path]
|
||||
if not os.path.exists(_)
|
||||
]
|
||||
|
||||
with open(args.instances_path, "r") as file:
|
||||
instance_list = json.load(file)
|
||||
config = Config.create(config_file=args.config_file)
|
||||
if not config.models:
|
||||
raise ValueError("No models found in config file.")
|
||||
if args.model_name not in config.models:
|
||||
raise ValueError(f"Model {args.model_name} not found in config file.")
|
||||
llm_config = config.models[args.model_name]
|
||||
llm_config.resolve_config_values()
|
||||
|
||||
candidate_dic = {}
|
||||
with open(args.candidate_path, "r") as file:
|
||||
for line in file.readlines():
|
||||
candidate = json.loads(line.strip())
|
||||
if "regressions" not in candidate:
|
||||
candidate["regressions"] = []
|
||||
for _ in range(len(candidate["patches"])):
|
||||
candidate["regressions"].append([])
|
||||
candidate_dic[candidate["instance_id"]] = candidate
|
||||
|
||||
tools_path = Path(__file__).parent / "trae_selector/tools"
|
||||
|
||||
try:
|
||||
log_path = Path(args.log_path)
|
||||
log_path.mkdir(parents=True, exist_ok=True)
|
||||
except Exception:
|
||||
print(f"Error creating log path for {args.log_path}")
|
||||
exit()
|
||||
|
||||
evaluation = SelectorEvaluation(
|
||||
llm_config,
|
||||
args.num_candidate,
|
||||
args.max_retry,
|
||||
args.max_turn,
|
||||
args.log_path,
|
||||
args.output_path,
|
||||
args.patches_path,
|
||||
instance_list,
|
||||
candidate_dic,
|
||||
tools_path.as_posix(),
|
||||
args.statistics_path,
|
||||
args.group_size,
|
||||
majority_voting=args.majority_voting,
|
||||
)
|
||||
|
||||
# evaluation.run_one("astropy__astropy-14369")
|
||||
evaluation.run_all(max_workers=args.max_workers)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1 @@
|
||||
# Package for trae selector components
|
||||
+123
@@ -0,0 +1,123 @@
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
import docker
|
||||
import pexpect
|
||||
|
||||
|
||||
class Sandbox:
|
||||
def __init__(self, namespace: str, name: str, tag: str, instance: dict, tools_path: str):
|
||||
self.namespace = namespace
|
||||
self.name = name
|
||||
self.tag = tag
|
||||
self.client = docker.from_env()
|
||||
self.commit_id = instance["base_commit"]
|
||||
self.instance_id = instance["instance_id"]
|
||||
self.container = None
|
||||
self.shell = None
|
||||
self.tools_path = tools_path
|
||||
|
||||
def get_project_path(self):
|
||||
project_path = self.container.exec_run("pwd").output.decode().strip()
|
||||
return project_path
|
||||
|
||||
def start_container(self):
|
||||
image = f"{self.namespace}/{self.name}:{self.tag}"
|
||||
host_path = "/tmp"
|
||||
container_path = "/tmp"
|
||||
self.container = self.client.containers.run(
|
||||
image,
|
||||
detach=True,
|
||||
tty=True,
|
||||
stdin_open=True,
|
||||
privileged=True,
|
||||
volumes={host_path: {"bind": container_path, "mode": "rw"}},
|
||||
)
|
||||
print(f"Container {self.container.short_id} started with image {image}")
|
||||
|
||||
cmd = f"chmod -R 777 {self.tools_path} && docker cp {self.tools_path} {self.container.name}:/home/swe-bench/"
|
||||
subprocess.run(cmd, check=True, shell=True)
|
||||
|
||||
checkout_res = self.container.exec_run(f"git checkout {self.commit_id}")
|
||||
print("checkout: ", checkout_res)
|
||||
|
||||
def start_shell(self):
|
||||
if self.container:
|
||||
if self.shell and self.shell.isalive():
|
||||
self.shell.close(force=True)
|
||||
command = f"docker exec -it {self.container.id} /bin/bash"
|
||||
self.shell = pexpect.spawn(command, maxread=200000)
|
||||
self.shell.expect([r"\$ ", r"# "], timeout=10)
|
||||
else:
|
||||
raise Exception("Container not started. Call start_container() first.")
|
||||
|
||||
def get_session(self):
|
||||
self.start_shell()
|
||||
|
||||
class Session:
|
||||
def __init__(self, sandbox):
|
||||
self.sandbox = sandbox
|
||||
|
||||
def execute(self, command, timeout=60):
|
||||
try:
|
||||
if command[-1] != "&":
|
||||
self.sandbox.shell.sendline(command + " && sleep 0.5")
|
||||
else:
|
||||
self.sandbox.shell.sendline(command)
|
||||
self.sandbox.shell.before = b""
|
||||
self.sandbox.shell.after = b""
|
||||
self.sandbox.shell.buffer = b""
|
||||
time.sleep(2)
|
||||
self.sandbox.shell.expect([r"swe-bench@.*:.*\$ ", r"root@.*:.*# "], 60)
|
||||
try:
|
||||
output = (
|
||||
self.sandbox.shell.before.decode("utf-8")
|
||||
+ self.sandbox.shell.after.decode("utf-8")
|
||||
+ self.sandbox.shell.buffer.decode("utf-8")
|
||||
)
|
||||
except Exception:
|
||||
output = (
|
||||
self.sandbox.shell.before.decode("utf-8", errors="replace")
|
||||
+ self.sandbox.shell.after.decode("utf-8", errors="replace")
|
||||
+ self.sandbox.shell.buffer.decode("utf-8", errors="replace")
|
||||
)
|
||||
output_lines = output.split("\r\n")
|
||||
if len(output_lines) > 1:
|
||||
output_lines = output_lines[1:-1]
|
||||
result_message = "\n".join(output_lines).replace("\x1b[?2004l\r", "")
|
||||
return result_message
|
||||
except pexpect.TIMEOUT:
|
||||
partial_output = ""
|
||||
if isinstance(self.sandbox.shell.before, bytes):
|
||||
partial_output += self.sandbox.shell.before.decode("utf-8")
|
||||
if isinstance(self.sandbox.shell.after, bytes):
|
||||
partial_output += self.sandbox.shell.after.decode("utf-8")
|
||||
if isinstance(self.sandbox.shell.buffer, bytes):
|
||||
partial_output += self.sandbox.shell.buffer.decode("utf-8")
|
||||
partial_output_lines = partial_output.split("\n")
|
||||
if len(partial_output_lines) > 1:
|
||||
partial_output_lines = partial_output_lines[1:-1]
|
||||
partial_output = "\n".join(partial_output_lines)
|
||||
return (
|
||||
"### Observation: "
|
||||
+ f"Error: Command '{command}' timed out after {timeout} seconds. Partial output:\n + {partial_output}"
|
||||
)
|
||||
|
||||
def close(self):
|
||||
if self.sandbox.shell:
|
||||
self.sandbox.shell.sendline("exit")
|
||||
self.sandbox.shell.expect(pexpect.EOF)
|
||||
self.sandbox.shell.close(force=True)
|
||||
self.sandbox.shell = None
|
||||
|
||||
return Session(self)
|
||||
|
||||
def stop_container(self):
|
||||
if self.container:
|
||||
if self.shell and self.shell.isalive():
|
||||
self.shell.close(force=True)
|
||||
self.shell = None
|
||||
self.container.stop()
|
||||
self.container.remove()
|
||||
print(f"Container {self.container.short_id} stopped and removed")
|
||||
self.container = None
|
||||
+242
@@ -0,0 +1,242 @@
|
||||
import re
|
||||
import shlex
|
||||
|
||||
from trae_agent.tools import tools_registry
|
||||
from trae_agent.tools.base import Tool, ToolResult
|
||||
from trae_agent.utils.config import ModelConfig
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage, LLMResponse
|
||||
from trae_agent.utils.llm_clients.llm_client import LLMClient
|
||||
from trae_agent.utils.trajectory_recorder import TrajectoryRecorder
|
||||
|
||||
from .sandbox import Sandbox
|
||||
|
||||
|
||||
class CandidatePatch:
|
||||
def __init__(self, id, patch, cleaned_patch, is_success_regression, is_success_patch):
|
||||
self.id = id
|
||||
self.patch = patch
|
||||
self.cleaned_patch = cleaned_patch
|
||||
self.is_success_regression = is_success_regression
|
||||
self.is_success_patch = is_success_patch
|
||||
|
||||
|
||||
def build_system_prompt(candidate_length: int) -> str:
|
||||
init_prompt = f"""\
|
||||
# ROLE: Act as an expert code evaluator. Given a codebase, an github issue and **{candidate_length} candidate patches** proposed by your colleagues, your responsibility is to **select the correct one** to solve the issue.
|
||||
|
||||
# WORK PROCESS:
|
||||
You are given a software issue and multiple candidate patches. Your goal is to identify the patch that correctly resolves the issue.
|
||||
|
||||
Follow these steps methodically:
|
||||
|
||||
**1. Understand the Issue and Codebase**
|
||||
Carefully read the issue description to comprehend the problem. You may need to examine the codebase for context, including:
|
||||
(1) Code referenced in the issue description;
|
||||
(2) The original code modified by each patch;
|
||||
(3) Unchanged parts of the same file;
|
||||
(4) Related files, functions, or modules that interact with the affected code.
|
||||
|
||||
**2. Analyze the Candidate Patches**
|
||||
For each patch, analyze its logic and intended fix. Consider whether the changes align with the issue description and coding conventions.
|
||||
|
||||
**3. Validate Functionality (Optional but Recommended)**
|
||||
If needed, write and run unit tests to evaluate the correctness and potential side effects of each patch.
|
||||
|
||||
**4. Select the Best Patch**
|
||||
Choose the patch that best resolves the issue with minimal risk of introducing new problems.
|
||||
|
||||
# FINAL REPORT: If you have successfully selected the correct patch, submit your answer in the following format:
|
||||
### Status: succeed
|
||||
### Result: Patch-x
|
||||
### Analysis: [Explain why Patch-x is correct.]
|
||||
|
||||
# IMPORTANT TIPS:
|
||||
1. Never avoid making a selection.
|
||||
2. Do not propose new patches.
|
||||
3. There must be at least one correct patch.
|
||||
"""
|
||||
return init_prompt
|
||||
|
||||
|
||||
def parse_tool_response(answer: LLMResponse, finish_reason: str, sandbox_session):
|
||||
result: list[LLMMessage] = []
|
||||
print("finish_reason:", finish_reason)
|
||||
if answer.tool_calls and len(answer.tool_calls) > 0:
|
||||
for tool_call in answer.tool_calls:
|
||||
tool_call_id = tool_call.call_id
|
||||
tool_name = tool_call.name
|
||||
|
||||
if tool_name == "str_replace_based_edit_tool":
|
||||
cmd = "cd /home/swe-bench/tools/ && /home/swe-bench/py312/bin/python3 execute_str_replace_editor.py"
|
||||
elif tool_name == "bash":
|
||||
cmd = (
|
||||
"cd /home/swe-bench/tools/ && /home/swe-bench/py312/bin/python3 execute_bash.py"
|
||||
)
|
||||
else:
|
||||
tool_message = LLMMessage(
|
||||
role="user",
|
||||
content="The tool name you provided is not in the list. Please choose one from `str_replace_editor` or `bash`!",
|
||||
tool_result=ToolResult(
|
||||
call_id=tool_call_id,
|
||||
name=tool_name,
|
||||
success=False,
|
||||
error="The tool name you provided is not in the list. Please choose one from `str_replace_editor` or `bash`!",
|
||||
),
|
||||
)
|
||||
result.append(tool_message)
|
||||
continue
|
||||
|
||||
all_arguments_valid = True
|
||||
tool_arguments = tool_call.arguments
|
||||
for key in tool_arguments:
|
||||
if isinstance(tool_arguments[key], list):
|
||||
try:
|
||||
tool_arguments[key] = str([int(factor) for factor in tool_arguments[key]])
|
||||
cmd += f" --{key} {shlex.quote(tool_arguments[key])}"
|
||||
except Exception:
|
||||
pass
|
||||
elif isinstance(tool_arguments[key], (int, bool)):
|
||||
cmd += f" --{key} {tool_arguments[key]}"
|
||||
elif isinstance(tool_arguments[key], dict):
|
||||
all_arguments_valid = False
|
||||
break
|
||||
else:
|
||||
cmd += f" --{key} {shlex.quote(tool_arguments[key])}"
|
||||
|
||||
if not all_arguments_valid:
|
||||
print("Tool Call Status: -1")
|
||||
tool_message = LLMMessage(
|
||||
role="user",
|
||||
content="Failed call tool. One of the arguments is dict type, you need to check the definition the tool.",
|
||||
tool_result=ToolResult(
|
||||
call_id=tool_call_id,
|
||||
name=tool_name,
|
||||
success=False,
|
||||
error="Failed call tool. One of the arguments is dict type, you need to check the definition the tool.",
|
||||
),
|
||||
)
|
||||
result.append(tool_message)
|
||||
continue
|
||||
|
||||
cmd += " > /home/swe-bench/tools/log.out 2>&1"
|
||||
print(repr(cmd))
|
||||
_ = sandbox_session.execute(cmd)
|
||||
sandbox_res = sandbox_session.execute("cat /home/swe-bench/tools/log.out")
|
||||
status = ""
|
||||
status_line_index = -1
|
||||
sandbox_res_str_list = sandbox_res.split("\n")
|
||||
for index, line in enumerate(sandbox_res_str_list):
|
||||
if line.strip().startswith("Tool Call Status:"):
|
||||
status = line
|
||||
status_line_index = index
|
||||
break
|
||||
if status_line_index != -1:
|
||||
sandbox_res_str_list.pop(status_line_index)
|
||||
res_content = "\n".join(sandbox_res_str_list)
|
||||
print(status)
|
||||
tool_message = LLMMessage(
|
||||
role="user",
|
||||
content=res_content,
|
||||
tool_result=ToolResult(
|
||||
call_id=tool_call_id,
|
||||
name=tool_name,
|
||||
success=status != "Tool Call Status: -1",
|
||||
result=res_content,
|
||||
error=None if status != "Tool Call Status: -1" else res_content,
|
||||
),
|
||||
)
|
||||
result.append(tool_message)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class SelectorAgent:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
llm_config: ModelConfig,
|
||||
sandbox: Sandbox,
|
||||
project_path: str,
|
||||
issue_description: str,
|
||||
trajectory_file_name: str,
|
||||
candidate_list: list[CandidatePatch],
|
||||
max_turn: int = 50,
|
||||
):
|
||||
self.llm_config = llm_config
|
||||
self.max_turn = max_turn
|
||||
self.sandbox = sandbox
|
||||
self.sandbox_session = self.sandbox.get_session()
|
||||
self.sandbox_session.execute("git reset --hard HEAD")
|
||||
self.initial_messages: list[LLMMessage] = []
|
||||
self.candidate_list: list[CandidatePatch] = candidate_list
|
||||
self.project_path: str = project_path
|
||||
self.issue_description: str = issue_description
|
||||
self.tools: list[Tool] = [
|
||||
tools_registry[tool_name](model_provider=llm_config.model_provider.provider)
|
||||
for tool_name in ["bash", "str_replace_based_edit_tool"]
|
||||
]
|
||||
self.llm_client = LLMClient(llm_config)
|
||||
self.trajectory_recorder: TrajectoryRecorder = TrajectoryRecorder(trajectory_file_name)
|
||||
|
||||
self.initial_messages.append(
|
||||
LLMMessage(role="system", content=build_system_prompt(len(candidate_list)))
|
||||
)
|
||||
user_prompt = f"\n[Codebase path]:\n{project_path}\n\n[Github issue description]:\n```\n{issue_description}\n```\n\n[Candidate Patches]:"
|
||||
for idx in range(0, len(candidate_list)):
|
||||
user_prompt += f"\nPatch-{idx + 1}:\n```\n{candidate_list[idx].patch}\n```"
|
||||
user_message = LLMMessage(role="user", content=user_prompt)
|
||||
self.initial_messages.append(user_message)
|
||||
|
||||
def run(self):
|
||||
print(f"max_turn: {self.max_turn}")
|
||||
print(f"### User Prompt:\n{self.initial_messages[1].content}\n")
|
||||
|
||||
turn = 0
|
||||
final_id, final_patch = self.candidate_list[0].id, self.candidate_list[0].patch
|
||||
messages = self.initial_messages
|
||||
while turn < self.max_turn:
|
||||
turn += 1
|
||||
llm_response = self.llm_client.chat(messages, self.llm_config, self.tools)
|
||||
self.trajectory_recorder.record_llm_interaction(
|
||||
messages,
|
||||
llm_response,
|
||||
self.llm_config.model_provider.provider,
|
||||
self.llm_config.model,
|
||||
self.tools,
|
||||
)
|
||||
answer_content = llm_response.content
|
||||
print(f"\n### Selector's Answer({turn})\n", answer_content)
|
||||
messages: list[LLMMessage] = []
|
||||
match = re.search(
|
||||
r"(?:###\s*)?Status:\s*(success|succeed|successfully|successful)\s*\n\s*(?:###\s*)?Result:",
|
||||
answer_content,
|
||||
)
|
||||
|
||||
if match:
|
||||
print("Match-1:", match.group(1).strip())
|
||||
match = re.search(
|
||||
r"(?:###\s*)?Result:\s*(.+?)\s*(?:###\s*)?Analysis:", answer_content
|
||||
)
|
||||
if match:
|
||||
result = match.group(1).strip().split("Patch-")[-1]
|
||||
print("Match-2:", result)
|
||||
if result in [str(_ + 1) for _ in range(len(self.candidate_list))]:
|
||||
final_id = self.candidate_list[int(result) - 1].id
|
||||
final_patch = self.candidate_list[int(result) - 1].patch
|
||||
else:
|
||||
final_id = self.candidate_list[0].id
|
||||
final_patch = self.candidate_list[0].patch
|
||||
break
|
||||
else:
|
||||
messages += parse_tool_response(
|
||||
llm_response, llm_response.finish_reason or "", self.sandbox_session
|
||||
)
|
||||
if messages[-1].content and " seconds. Partial output:" in messages[-1].content:
|
||||
self.sandbox_session = self.sandbox.get_session()
|
||||
|
||||
print(f"\n### System Response({turn})\n", messages)
|
||||
self.trajectory_recorder.finalize_recording(True, final_patch)
|
||||
self.sandbox_session.execute("git reset --hard HEAD")
|
||||
self.sandbox_session.close()
|
||||
|
||||
return final_id, final_patch
|
||||
@@ -0,0 +1,389 @@
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
from collections import Counter
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from trae_agent.utils.config import ModelConfig
|
||||
|
||||
from .sandbox import Sandbox
|
||||
from .selector_agent import CandidatePatch, SelectorAgent
|
||||
from .utils import clean_patch, get_trajectory_filename, save_patches, save_selection_success
|
||||
|
||||
|
||||
def run_instance(
|
||||
*,
|
||||
instance,
|
||||
candidate_log,
|
||||
output_path,
|
||||
max_retry,
|
||||
num_candidate,
|
||||
tools_path,
|
||||
statistics_path,
|
||||
group_size,
|
||||
llm_config,
|
||||
max_turn,
|
||||
log_path,
|
||||
patches_path,
|
||||
majority_voting=True,
|
||||
):
|
||||
# candidate_log is a list of num_candidate candidate patches
|
||||
# divide candidate_log into groups of group_size
|
||||
groups = []
|
||||
for i in range(0, num_candidate, group_size):
|
||||
this_group = {
|
||||
"instance_id": candidate_log["instance_id"],
|
||||
"issue": candidate_log["issue"],
|
||||
"patches": candidate_log["patches"][i : i + group_size],
|
||||
"regressions": candidate_log["regressions"][i : i + group_size],
|
||||
"success_id": candidate_log["success_id"][i : i + group_size],
|
||||
}
|
||||
groups.append(this_group)
|
||||
|
||||
for group_id, group in enumerate(groups):
|
||||
run_instance_by_group(
|
||||
instance=instance,
|
||||
candidate_log=group,
|
||||
output_path=output_path,
|
||||
max_retry=max_retry,
|
||||
num_candidate=len(group),
|
||||
tools_path=tools_path,
|
||||
statistics_path=statistics_path,
|
||||
llm_config=llm_config,
|
||||
max_turn=max_turn,
|
||||
log_path=log_path,
|
||||
patches_path=patches_path,
|
||||
group_id=group_id,
|
||||
num_groups=len(groups),
|
||||
majority_voting=majority_voting,
|
||||
)
|
||||
|
||||
|
||||
def run_instance_by_group(
|
||||
*,
|
||||
instance,
|
||||
candidate_log,
|
||||
output_path,
|
||||
max_retry,
|
||||
num_candidate,
|
||||
tools_path,
|
||||
statistics_path,
|
||||
llm_config,
|
||||
max_turn,
|
||||
log_path,
|
||||
patches_path,
|
||||
group_id,
|
||||
num_groups,
|
||||
majority_voting=True,
|
||||
):
|
||||
print(f"[Group {group_id}/{num_groups}] processing: {instance['instance_id']}")
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
|
||||
# check if the group has already been processed: the statistics json file exists and is not empty
|
||||
file_path = statistics_path + f"/group_{group_id}/{instance['instance_id']}.json"
|
||||
if os.path.exists(file_path) and os.path.getsize(file_path) > 0:
|
||||
print(
|
||||
f"[Group {group_id}/{num_groups}] for instance {instance['instance_id']} has already been processed. Skipping..."
|
||||
)
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
sys.stdout = sys.__stdout__
|
||||
sys.stderr = sys.__stderr__
|
||||
return
|
||||
|
||||
# check if the group is all failed or all success. If so, skip this group
|
||||
all_failed = True
|
||||
all_success = True
|
||||
for success_id in candidate_log["success_id"]:
|
||||
if success_id == 1:
|
||||
all_failed = False
|
||||
if success_id != 1:
|
||||
all_success = False
|
||||
if all_failed or all_success:
|
||||
print(
|
||||
f"[Group ID {group_id} in {num_groups}] groups for instance {instance['instance_id']} {'all failed' if all_failed else 'all success'}. Skipping..."
|
||||
)
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
sys.stdout = sys.__stdout__
|
||||
sys.stderr = sys.__stderr__
|
||||
|
||||
save_patches(
|
||||
instance_id=instance["instance_id"],
|
||||
patches_path=patches_path,
|
||||
patches=candidate_log["patches"][0],
|
||||
group_id=group_id,
|
||||
)
|
||||
|
||||
if all_failed:
|
||||
save_selection_success(
|
||||
instance_id=instance["instance_id"],
|
||||
statistics_path=statistics_path,
|
||||
patch_id=0,
|
||||
is_success=0,
|
||||
group_id=group_id,
|
||||
is_all_failed=True,
|
||||
is_all_success=False,
|
||||
)
|
||||
if all_success:
|
||||
save_selection_success(
|
||||
instance_id=instance["instance_id"],
|
||||
statistics_path=statistics_path,
|
||||
patch_id=0,
|
||||
is_success=1,
|
||||
group_id=group_id,
|
||||
is_all_success=True,
|
||||
is_all_failed=False,
|
||||
)
|
||||
|
||||
return
|
||||
|
||||
log_dir_path = Path(output_path) / f"group_{group_id}"
|
||||
log_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
log_file_path = log_dir_path / f"{instance['instance_id']}.log"
|
||||
with open(log_file_path, "w") as log_file:
|
||||
sys.stdout = log_file
|
||||
sys.stderr = log_file
|
||||
namespace = "swebench"
|
||||
image_name = "sweb.eval.x86_64." + instance["instance_id"].replace("__", "_1776_")
|
||||
tag = "latest"
|
||||
|
||||
try:
|
||||
current_try = 0
|
||||
while current_try < max_retry:
|
||||
print("current_try:", current_try)
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
print("time: ", datetime.now().strftime("%Y%m%d%H%M%S"))
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
current_try += 1
|
||||
sandbox = None
|
||||
try:
|
||||
candidate_list = []
|
||||
for idx in range(len(candidate_log["patches"])):
|
||||
if candidate_log["patches"][idx].strip() == "":
|
||||
continue
|
||||
cleaned_patch = clean_patch(candidate_log["patches"][idx])
|
||||
is_success_regression = len(candidate_log["regressions"][idx]) == 0
|
||||
candidate_list.append(
|
||||
CandidatePatch(
|
||||
idx,
|
||||
candidate_log["patches"][idx],
|
||||
cleaned_patch,
|
||||
is_success_regression,
|
||||
candidate_log["success_id"][idx],
|
||||
)
|
||||
)
|
||||
|
||||
# regression testing
|
||||
candidate_list_regression = [
|
||||
candidate for candidate in candidate_list if candidate.is_success_regression
|
||||
]
|
||||
if len(candidate_list_regression):
|
||||
candidate_list = candidate_list_regression
|
||||
print(f"[Retry No:{current_try}] regression testing done")
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
|
||||
# patch deduplication
|
||||
candidate_list_deduplication, cleaned_candidate_set = [], set()
|
||||
for candidate in candidate_list:
|
||||
if candidate.cleaned_patch not in cleaned_candidate_set:
|
||||
cleaned_candidate_set.add(candidate.cleaned_patch)
|
||||
candidate_list_deduplication.append(candidate)
|
||||
candidate_list = candidate_list_deduplication
|
||||
print(f"[Retry No:{current_try}] patch deduplication done")
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
|
||||
# sandbox & tools
|
||||
sandbox = Sandbox(namespace, image_name, tag, instance, tools_path)
|
||||
sandbox.start_container()
|
||||
project_path = sandbox.get_project_path()
|
||||
print(f"[Retry No:{current_try}] sandbox & tools done")
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
|
||||
# majority voting
|
||||
if majority_voting:
|
||||
final_id_list, final_patch_list = [], []
|
||||
for idx in range(num_candidate):
|
||||
select_agent = SelectorAgent(
|
||||
llm_config=llm_config,
|
||||
sandbox=sandbox,
|
||||
project_path=project_path,
|
||||
issue_description=instance["problem_statement"],
|
||||
trajectory_file_name=get_trajectory_filename(
|
||||
instance["instance_id"], log_path, group_id, idx
|
||||
),
|
||||
candidate_list=candidate_list,
|
||||
max_turn=max_turn,
|
||||
)
|
||||
|
||||
final_id, final_patch = select_agent.run()
|
||||
final_id_list.append(final_id)
|
||||
final_patch_list.append(final_patch)
|
||||
if max(Counter(final_id_list).values()) > num_candidate / 2:
|
||||
break
|
||||
print(f"[Retry No:{current_try}] majority voting done")
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
|
||||
counter = Counter(final_id_list)
|
||||
max_count = max(counter.values())
|
||||
most_common_ids = [
|
||||
elem for elem, count in counter.items() if count == max_count
|
||||
]
|
||||
result = {}
|
||||
for id_ in most_common_ids:
|
||||
indexes = [i for i, val in enumerate(final_id_list) if val == id_]
|
||||
result[id_] = indexes
|
||||
final_id = most_common_ids[0]
|
||||
final_patch = final_patch_list[result[final_id][0]]
|
||||
print(f"[Retry No:{current_try}] final_id_list: {final_id_list}")
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
else:
|
||||
select_agent = SelectorAgent(
|
||||
llm_config=llm_config,
|
||||
sandbox=sandbox,
|
||||
project_path=project_path,
|
||||
issue_description=instance["problem_statement"],
|
||||
trajectory_file_name=get_trajectory_filename(
|
||||
instance["instance_id"], log_path, group_id, 0
|
||||
),
|
||||
candidate_list=candidate_list,
|
||||
max_turn=max_turn,
|
||||
)
|
||||
final_id, final_patch = select_agent.run()
|
||||
save_patches(
|
||||
instance_id=instance["instance_id"],
|
||||
patches_path=patches_path,
|
||||
patches=final_patch,
|
||||
group_id=group_id,
|
||||
)
|
||||
|
||||
is_success_patch = 0
|
||||
for candidate in candidate_list:
|
||||
if final_id == candidate.id:
|
||||
is_success_patch = candidate.is_success_patch
|
||||
save_selection_success(
|
||||
instance_id=instance["instance_id"],
|
||||
statistics_path=statistics_path,
|
||||
patch_id=final_id,
|
||||
is_success=is_success_patch,
|
||||
group_id=group_id,
|
||||
)
|
||||
sandbox.stop_container()
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"Error occurred: {e}")
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
print("Detailed Error:\n", traceback.format_exc())
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
if sandbox is not None:
|
||||
sandbox.stop_container()
|
||||
finally:
|
||||
sys.stdout = sys.__stdout__
|
||||
sys.stderr = sys.__stderr__
|
||||
print(f" finished: {instance['instance_id']}")
|
||||
|
||||
|
||||
class SelectorEvaluation:
|
||||
def __init__(
|
||||
self,
|
||||
llm_config: ModelConfig,
|
||||
num_candidate: int,
|
||||
max_retry: int,
|
||||
max_turn: int,
|
||||
log_path: str,
|
||||
output_path: str,
|
||||
patches_path: str,
|
||||
instance_list: list,
|
||||
candidate_dic: dict[str, dict],
|
||||
tools_path: str,
|
||||
statistics_path: str,
|
||||
group_size: int,
|
||||
majority_voting: bool = True,
|
||||
):
|
||||
self.llm_config = llm_config
|
||||
self.num_candidate = num_candidate
|
||||
self.max_retry = max_retry
|
||||
self.log_path = log_path
|
||||
self.output_path = output_path
|
||||
self.patches_path = patches_path
|
||||
self.instance_list = instance_list
|
||||
self.candidate_dic = candidate_dic
|
||||
self.max_turn = max_turn
|
||||
self.tools_path = tools_path
|
||||
self.statistics_path = statistics_path
|
||||
self.group_size = group_size
|
||||
self.majority_voting = majority_voting
|
||||
|
||||
def run_all(self, max_workers=None):
|
||||
"""Run all instances concurrently using ThreadPoolExecutor.
|
||||
|
||||
Args:
|
||||
max_workers: Maximum number of worker threads. If None, defaults to min(32, os.cpu_count() + 4)
|
||||
"""
|
||||
with ProcessPoolExecutor(max_workers=max_workers) as ex:
|
||||
futures = {
|
||||
ex.submit(
|
||||
run_instance,
|
||||
instance=instance,
|
||||
candidate_log=self.candidate_dic[instance["instance_id"]],
|
||||
output_path=self.output_path,
|
||||
max_retry=self.max_retry,
|
||||
num_candidate=self.num_candidate,
|
||||
tools_path=self.tools_path,
|
||||
statistics_path=self.statistics_path,
|
||||
group_size=self.group_size,
|
||||
llm_config=self.llm_config,
|
||||
max_turn=self.max_turn,
|
||||
log_path=self.log_path,
|
||||
patches_path=self.patches_path,
|
||||
majority_voting=self.majority_voting,
|
||||
): instance["instance_id"]
|
||||
for instance in self.instance_list
|
||||
}
|
||||
|
||||
with tqdm(total=len(futures), ascii=True, desc="Processing instances") as pbar:
|
||||
for fut in as_completed(futures):
|
||||
iid = futures[fut]
|
||||
try:
|
||||
result_iid = fut.result()
|
||||
pbar.set_postfix({"completed": result_iid})
|
||||
except Exception:
|
||||
result_iid = iid
|
||||
print(traceback.format_exc())
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
finally:
|
||||
pbar.update(1)
|
||||
|
||||
def run_one(self, instance_id):
|
||||
for idx in range(len(self.instance_list)):
|
||||
if instance_id == self.instance_list[idx]["instance_id"]:
|
||||
run_instance(
|
||||
instance=self.instance_list[idx],
|
||||
candidate_log=self.candidate_dic[instance_id],
|
||||
output_path=self.output_path,
|
||||
max_retry=self.max_retry,
|
||||
num_candidate=self.num_candidate,
|
||||
tools_path=self.tools_path,
|
||||
statistics_path=self.statistics_path,
|
||||
group_size=self.group_size,
|
||||
llm_config=self.llm_config,
|
||||
max_turn=self.max_turn,
|
||||
log_path=self.log_path,
|
||||
patches_path=self.patches_path,
|
||||
majority_voting=self.majority_voting,
|
||||
)
|
||||
@@ -0,0 +1,46 @@
|
||||
from dataclasses import dataclass, fields, replace
|
||||
|
||||
|
||||
@dataclass(kw_only=True, frozen=True)
|
||||
class ToolResult:
|
||||
output: str | None = None
|
||||
error: str | None = None
|
||||
base64_image: str | None = None
|
||||
system: str | None = None
|
||||
|
||||
def __bool__(self):
|
||||
return any(getattr(self, field.name) for field in fields(self))
|
||||
|
||||
def __add__(self, other: "ToolResult"):
|
||||
def combine_fields(field: str | None, other_field: str | None, concatenate: bool = True):
|
||||
if field and other_field:
|
||||
if concatenate:
|
||||
return field + other_field
|
||||
raise ValueError("Cannot combine tool results")
|
||||
return field or other_field
|
||||
|
||||
return ToolResult(
|
||||
output=combine_fields(self.output, other.output),
|
||||
error=combine_fields(self.error, other.error),
|
||||
base64_image=combine_fields(self.base64_image, other.base64_image, False),
|
||||
system=combine_fields(self.system, other.system),
|
||||
)
|
||||
|
||||
def replace(self, **kwargs):
|
||||
return replace(self, **kwargs)
|
||||
|
||||
|
||||
class CLIResult(ToolResult):
|
||||
"""A ToolResult that can be rendered as a CLI output."""
|
||||
|
||||
|
||||
class ToolFailure(ToolResult):
|
||||
"""A ToolResult that represents a failure."""
|
||||
|
||||
|
||||
class ToolError(Exception):
|
||||
"""Raised when a tool encounters an error."""
|
||||
|
||||
def __init__(self, message: str):
|
||||
super().__init__(message)
|
||||
self.message: str = message
|
||||
@@ -0,0 +1,122 @@
|
||||
import asyncio
|
||||
import os
|
||||
from typing import ClassVar, Literal
|
||||
|
||||
from base import CLIResult, ToolError, ToolResult
|
||||
|
||||
|
||||
class _BashSession:
|
||||
_started: bool
|
||||
_process: asyncio.subprocess.Process
|
||||
|
||||
command: str = "/bin/bash"
|
||||
_output_delay: float = 0.2
|
||||
_timeout: float = 120.0
|
||||
_sentinel: str = "<<exit>>"
|
||||
|
||||
def __init__(self):
|
||||
self._started = False
|
||||
self._timed_out = False
|
||||
|
||||
async def start(self):
|
||||
if self._started:
|
||||
return
|
||||
|
||||
self._process = await asyncio.create_subprocess_shell(
|
||||
self.command,
|
||||
preexec_fn=os.setsid,
|
||||
shell=True,
|
||||
bufsize=0,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
|
||||
self._started = True
|
||||
|
||||
def stop(self):
|
||||
if not self._started:
|
||||
raise ToolError("Session has not started.")
|
||||
if self._process.returncode is not None:
|
||||
return
|
||||
self._process.terminate()
|
||||
|
||||
async def run(self, command: str):
|
||||
if not self._started:
|
||||
raise ToolError("Session has not started.")
|
||||
if self._process.returncode is not None:
|
||||
return ToolResult(
|
||||
system="tool must be restarted",
|
||||
error=f"bash has exited with returncode {self._process.returncode}",
|
||||
)
|
||||
if self._timed_out:
|
||||
raise ToolError(
|
||||
f"timed out: bash has not returned in {self._timeout} seconds and must be restarted",
|
||||
)
|
||||
|
||||
assert self._process.stdin
|
||||
assert self._process.stdout
|
||||
assert self._process.stderr
|
||||
|
||||
self._process.stdin.write(command.encode() + f"; echo '{self._sentinel}'\n".encode())
|
||||
await self._process.stdin.drain()
|
||||
|
||||
try:
|
||||
async with asyncio.timeout(self._timeout):
|
||||
while True:
|
||||
await asyncio.sleep(self._output_delay)
|
||||
output = self._process.stdout._buffer.decode()
|
||||
if self._sentinel in output:
|
||||
output = output[: output.index(self._sentinel)]
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
self._timed_out = True
|
||||
raise ToolError(
|
||||
f"timed out: bash has not returned in {self._timeout} seconds and must be restarted",
|
||||
) from None
|
||||
|
||||
if output.endswith("\n"):
|
||||
output = output[:-1]
|
||||
|
||||
error = self._process.stderr._buffer.decode()
|
||||
if error.endswith("\n"):
|
||||
error = error[:-1]
|
||||
|
||||
self._process.stdout._buffer.clear()
|
||||
self._process.stderr._buffer.clear()
|
||||
|
||||
return CLIResult(output=output, error=error)
|
||||
|
||||
|
||||
class BashTool:
|
||||
_session: _BashSession | None
|
||||
name: ClassVar[Literal["bash"]] = "bash"
|
||||
api_type: ClassVar[Literal["bash_2025"]] = "bash_2025"
|
||||
|
||||
def __init__(self):
|
||||
self._session = None
|
||||
super().__init__()
|
||||
|
||||
async def __call__(self, command: str | None = None, restart: bool = False, **kwargs):
|
||||
if restart:
|
||||
if self._session:
|
||||
self._session.stop()
|
||||
self._session = _BashSession()
|
||||
await self._session.start()
|
||||
|
||||
return ToolResult(system="tool has been restarted.")
|
||||
|
||||
if self._session is None:
|
||||
self._session = _BashSession()
|
||||
await self._session.start()
|
||||
|
||||
if command is not None:
|
||||
return await self._session.run(command)
|
||||
|
||||
raise ToolError("no command provided.")
|
||||
|
||||
def to_params(self):
|
||||
return {
|
||||
"type": self.api_type,
|
||||
"name": self.name,
|
||||
}
|
||||
@@ -0,0 +1,244 @@
|
||||
import os
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Literal, get_args
|
||||
|
||||
from base import CLIResult, ToolError, ToolResult
|
||||
from run import maybe_truncate, run
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
|
||||
Command = Literal[
|
||||
"view",
|
||||
"create",
|
||||
"str_replace",
|
||||
"insert",
|
||||
"undo_edit",
|
||||
]
|
||||
SNIPPET_LINES: int = 4
|
||||
|
||||
|
||||
def write_text(filename, content):
|
||||
with open(str(filename), "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
|
||||
|
||||
class EditTool:
|
||||
api_type: Literal["text_editor_2025"] = "text_editor_2025"
|
||||
name: Literal["str_replace_editor"] = "str_replace_editor"
|
||||
|
||||
_file_history: dict[Path, list[str]]
|
||||
|
||||
def __init__(self):
|
||||
self._file_history = defaultdict(list)
|
||||
super().__init__()
|
||||
|
||||
def to_params(self):
|
||||
return {
|
||||
"name": self.name,
|
||||
"type": self.api_type,
|
||||
}
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
*,
|
||||
command: Command,
|
||||
path: str,
|
||||
file_text: str | None = None,
|
||||
view_range: list[int] | None = None,
|
||||
old_str: str | None = None,
|
||||
new_str: str | None = None,
|
||||
insert_line: int | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
_path = Path(path)
|
||||
self.validate_path(command, _path)
|
||||
if command == "view":
|
||||
return await self.view(_path, view_range)
|
||||
elif command == "create":
|
||||
if file_text is None:
|
||||
raise ToolError("Parameter `file_text` is required for command: create")
|
||||
self.write_file(_path, file_text)
|
||||
self._file_history[_path].append(file_text)
|
||||
return ToolResult(output=f"File created successfully at: {_path}")
|
||||
elif command == "str_replace":
|
||||
if old_str is None:
|
||||
raise ToolError("Parameter `old_str` is required for command: str_replace")
|
||||
return self.str_replace(_path, old_str, new_str)
|
||||
elif command == "insert":
|
||||
if insert_line is None:
|
||||
raise ToolError("Parameter `insert_line` is required for command: insert")
|
||||
if new_str is None:
|
||||
raise ToolError("Parameter `new_str` is required for command: insert")
|
||||
return self.insert(_path, insert_line, new_str)
|
||||
elif command == "undo_edit":
|
||||
return self.undo_edit(_path)
|
||||
raise ToolError(
|
||||
f"Unrecognized command {command}. The allowed commands for the {self.name} tool are: {', '.join(get_args(Command))}"
|
||||
)
|
||||
|
||||
def validate_path(self, command: str, path: Path):
|
||||
if not path.is_absolute():
|
||||
suggested_path = Path("") / path
|
||||
raise ToolError(
|
||||
f"The path {path} is not an absolute path, it should start with `/`. Maybe you meant {suggested_path}?"
|
||||
)
|
||||
if not path.exists() and command != "create":
|
||||
raise ToolError(f"The path {path} does not exist. Please provide a valid path.")
|
||||
if path.exists() and command == "create":
|
||||
raise ToolError(
|
||||
f"File already exists at: {path}. Cannot overwrite files using command `create`."
|
||||
)
|
||||
if path.is_dir() and command != "view":
|
||||
raise ToolError(
|
||||
f"The path {path} is a directory and only the `view` command can be used on directories"
|
||||
)
|
||||
|
||||
async def view(self, path: Path, view_range: list[int] | None = None):
|
||||
if path.is_dir():
|
||||
if view_range:
|
||||
raise ToolError(
|
||||
"The `view_range` parameter is not allowed when `path` points to a directory."
|
||||
)
|
||||
|
||||
_, stdout, stderr = await run(rf"find {path} -maxdepth 2 -not -path '*/\.*'")
|
||||
if not stderr:
|
||||
stdout = f"Here's the files and directories up to 2 levels deep in {path}, excluding hidden items:\n{stdout}\n"
|
||||
return CLIResult(output=stdout, error=stderr)
|
||||
|
||||
file_content = self.read_file(path)
|
||||
init_line = 1
|
||||
if view_range:
|
||||
if len(view_range) != 2 or not all(isinstance(i, int) for i in view_range):
|
||||
raise ToolError("Invalid `view_range`. It should be a list of two integers.")
|
||||
file_lines = file_content.split("\n")
|
||||
n_lines_file = len(file_lines)
|
||||
init_line, final_line = view_range
|
||||
if init_line < 1 or init_line > n_lines_file:
|
||||
raise ToolError(
|
||||
f"Invalid `view_range`: {view_range}. Its first element `{init_line}` should be within the range of lines of the file: {[1, n_lines_file]}"
|
||||
)
|
||||
if final_line > n_lines_file:
|
||||
raise ToolError(
|
||||
f"Invalid `view_range`: {view_range}. Its second element `{final_line}` should be smaller than the number of lines in the file: `{n_lines_file}`"
|
||||
)
|
||||
if final_line != -1 and final_line < init_line:
|
||||
raise ToolError(
|
||||
f"Invalid `view_range`: {view_range}. Its second element `{final_line}` should be larger or equal than its first `{init_line}`"
|
||||
)
|
||||
|
||||
if final_line == -1:
|
||||
file_content = "\n".join(file_lines[init_line - 1 :])
|
||||
else:
|
||||
file_content = "\n".join(file_lines[init_line - 1 : final_line])
|
||||
|
||||
return CLIResult(output=self._make_output(file_content, str(path), init_line=init_line))
|
||||
|
||||
def str_replace(self, path: Path, old_str: str, new_str: str | None):
|
||||
file_content = self.read_file(path).expandtabs()
|
||||
old_str = old_str.expandtabs()
|
||||
new_str = new_str.expandtabs() if new_str is not None else ""
|
||||
|
||||
occurrences = file_content.count(old_str)
|
||||
if occurrences == 0:
|
||||
raise ToolError(
|
||||
f"No replacement was performed, old_str `{old_str}` did not appear verbatim in {path}."
|
||||
)
|
||||
elif occurrences > 1:
|
||||
file_content_lines = file_content.split("\n")
|
||||
lines = [idx + 1 for idx, line in enumerate(file_content_lines) if old_str in line]
|
||||
raise ToolError(
|
||||
f"No replacement was performed. Multiple occurrences of old_str `{old_str}` in lines {lines}. Please ensure it is unique"
|
||||
)
|
||||
|
||||
new_file_content = file_content.replace(old_str, new_str)
|
||||
|
||||
self.write_file(path, new_file_content)
|
||||
self._file_history[path].append(file_content)
|
||||
|
||||
replacement_line = file_content.split(old_str)[0].count("\n")
|
||||
start_line = max(0, replacement_line - SNIPPET_LINES)
|
||||
end_line = replacement_line + SNIPPET_LINES + new_str.count("\n")
|
||||
snippet = "\n".join(new_file_content.split("\n")[start_line : end_line + 1])
|
||||
|
||||
success_msg = f"The file {path} has been edited. "
|
||||
success_msg += self._make_output(snippet, f"a snippet of {path}", start_line + 1)
|
||||
success_msg += "Review the changes and make sure they are as expected. Edit the file again if necessary."
|
||||
|
||||
return CLIResult(output=success_msg)
|
||||
|
||||
def insert(self, path: Path, insert_line: int, new_str: str):
|
||||
file_text = self.read_file(path).expandtabs()
|
||||
new_str = new_str.expandtabs()
|
||||
file_text_lines = file_text.split("\n")
|
||||
n_lines_file = len(file_text_lines)
|
||||
|
||||
if insert_line < 0 or insert_line > n_lines_file:
|
||||
raise ToolError(
|
||||
f"Invalid `insert_line` parameter: {insert_line}. It should be within the range of lines of the file: {[0, n_lines_file]}"
|
||||
)
|
||||
|
||||
new_str_lines = new_str.split("\n")
|
||||
new_file_text_lines = (
|
||||
file_text_lines[:insert_line] + new_str_lines + file_text_lines[insert_line:]
|
||||
)
|
||||
snippet_lines = (
|
||||
file_text_lines[max(0, insert_line - SNIPPET_LINES) : insert_line]
|
||||
+ new_str_lines
|
||||
+ file_text_lines[insert_line : insert_line + SNIPPET_LINES]
|
||||
)
|
||||
|
||||
new_file_text = "\n".join(new_file_text_lines)
|
||||
snippet = "\n".join(snippet_lines)
|
||||
|
||||
self.write_file(path, new_file_text)
|
||||
self._file_history[path].append(file_text)
|
||||
|
||||
success_msg = f"The file {path} has been edited. "
|
||||
success_msg += self._make_output(
|
||||
snippet,
|
||||
"a snippet of the edited file",
|
||||
max(1, insert_line - SNIPPET_LINES + 1),
|
||||
)
|
||||
success_msg += "Review the changes and make sure they are as expected (correct indentation, no duplicate lines, etc). Edit the file again if necessary."
|
||||
return CLIResult(output=success_msg)
|
||||
|
||||
def undo_edit(self, path: Path):
|
||||
if not self._file_history[path]:
|
||||
raise ToolError(f"No edit history found for {path}.")
|
||||
|
||||
old_text = self._file_history[path].pop()
|
||||
self.write_file(path, old_text)
|
||||
|
||||
return CLIResult(
|
||||
output=f"Last edit to {path} undone successfully. {self._make_output(old_text, str(path))}"
|
||||
)
|
||||
|
||||
def read_file(self, path: Path):
|
||||
try:
|
||||
return path.read_text()
|
||||
except Exception as e:
|
||||
raise ToolError(f"Ran into {e} while trying to read {path}") from None
|
||||
|
||||
def write_file(self, path: Path, file: str):
|
||||
try:
|
||||
path.write_text(file)
|
||||
except Exception as e:
|
||||
raise ToolError(f"Ran into {e} while trying to write to {path}") from None
|
||||
|
||||
def _make_output(
|
||||
self,
|
||||
file_content: str,
|
||||
file_descriptor: str,
|
||||
init_line: int = 1,
|
||||
expand_tabs: bool = True,
|
||||
):
|
||||
file_content = maybe_truncate(file_content)
|
||||
if expand_tabs:
|
||||
file_content = file_content.expandtabs()
|
||||
file_content = "\n".join(
|
||||
[f"{i + init_line:6}\t{line}" for i, line in enumerate(file_content.split("\n"))]
|
||||
)
|
||||
return (
|
||||
f"Here's the result of running `cat -n` on {file_descriptor}:\n" + file_content + "\n"
|
||||
)
|
||||
@@ -0,0 +1,44 @@
|
||||
import asyncio
|
||||
import sys
|
||||
|
||||
from base import ToolError
|
||||
from bash import BashTool
|
||||
|
||||
|
||||
async def execute_command(**kwargs):
|
||||
tool = BashTool()
|
||||
|
||||
if kwargs.get("restart") is None:
|
||||
kwargs["restart"] = False
|
||||
elif kwargs.get("restart").lower() == "true":
|
||||
kwargs["restart"] = True
|
||||
else:
|
||||
kwargs["restart"] = False
|
||||
|
||||
try:
|
||||
result = await tool(command=kwargs.get("command"), restart=kwargs.get("restart"))
|
||||
return_content = ""
|
||||
if result.output is not None:
|
||||
return_content += result.output
|
||||
if result.error is not None:
|
||||
return_content += "\n" + result.error
|
||||
return 0, return_content
|
||||
except ToolError as e:
|
||||
return -1, e
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = sys.argv[1:]
|
||||
kwargs = {}
|
||||
it = iter(args)
|
||||
for arg in it:
|
||||
if arg.startswith("--"):
|
||||
key = arg.lstrip("-")
|
||||
try:
|
||||
value = next(it)
|
||||
kwargs[key] = value
|
||||
except StopIteration:
|
||||
kwargs[key] = None
|
||||
status, output = asyncio.run(execute_command(**kwargs))
|
||||
print(f"Tool Call Status: {status}")
|
||||
print(output)
|
||||
+68
@@ -0,0 +1,68 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import pickle
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from base import ToolError
|
||||
from edit import EditTool
|
||||
|
||||
|
||||
async def execute_command(**kwargs):
|
||||
tool = EditTool()
|
||||
|
||||
if os.path.exists("file_history.pkl"):
|
||||
with open("file_history.pkl", "rb") as file:
|
||||
tool._file_history = pickle.load(file)
|
||||
|
||||
kwargs["path"] = Path(kwargs["path"]) if "path" in kwargs and kwargs["path"] else None
|
||||
|
||||
with contextlib.suppress(json.JSONDecodeError):
|
||||
kwargs["view_range"] = (
|
||||
json.loads(kwargs["view_range"]) if kwargs.get("view_range") is not None else None
|
||||
)
|
||||
|
||||
with contextlib.suppress(ValueError):
|
||||
kwargs["insert_line"] = (
|
||||
int(kwargs["insert_line"]) if kwargs.get("insert_line") is not None else None
|
||||
)
|
||||
|
||||
try:
|
||||
result = await tool(
|
||||
command=kwargs.get("command"),
|
||||
path=kwargs.get("path"),
|
||||
file_text=kwargs.get("file_text"),
|
||||
view_range=kwargs.get("view_range"),
|
||||
insert_line=kwargs.get("insert_line"),
|
||||
old_str=kwargs.get("old_str"),
|
||||
new_str=kwargs.get("new_str"),
|
||||
)
|
||||
with open("file_history.pkl", "wb") as file:
|
||||
pickle.dump(tool._file_history, file)
|
||||
return_content = ""
|
||||
if result.output is not None:
|
||||
return_content += result.output
|
||||
if result.error is not None:
|
||||
return_content += "\n" + result.error
|
||||
return 0, return_content
|
||||
except ToolError as e:
|
||||
return -1, e
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = sys.argv[1:]
|
||||
kwargs = {}
|
||||
it = iter(args)
|
||||
for arg in it:
|
||||
if arg.startswith("--"):
|
||||
key = arg.lstrip("-")
|
||||
try:
|
||||
value = next(it)
|
||||
kwargs[key] = value
|
||||
except StopIteration:
|
||||
kwargs[key] = None
|
||||
status, output = asyncio.run(execute_command(**kwargs))
|
||||
print(f"Tool Call Status: {status}")
|
||||
print(output)
|
||||
@@ -0,0 +1,36 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
|
||||
TRUNCATED_MESSAGE: str = "<response clipped><NOTE>To save on context only part of this file has been shown to you. You should retry this tool after you have searched inside the file with `grep -n` in order to find the line numbers of what you are looking for.</NOTE>"
|
||||
MAX_RESPONSE_LEN: int = 16000
|
||||
|
||||
|
||||
def maybe_truncate(content: str, truncate_after: int | None = MAX_RESPONSE_LEN):
|
||||
return (
|
||||
content
|
||||
if not truncate_after or len(content) <= truncate_after
|
||||
else content[:truncate_after] + TRUNCATED_MESSAGE
|
||||
)
|
||||
|
||||
|
||||
async def run(
|
||||
cmd: str,
|
||||
timeout: float | None = 120.0,
|
||||
truncate_after: int | None = MAX_RESPONSE_LEN,
|
||||
):
|
||||
process = await asyncio.create_subprocess_shell(
|
||||
cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
|
||||
try:
|
||||
stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout)
|
||||
return (
|
||||
process.returncode or 0,
|
||||
maybe_truncate(stdout.decode(), truncate_after=truncate_after),
|
||||
maybe_truncate(stderr.decode(), truncate_after=truncate_after),
|
||||
)
|
||||
except asyncio.TimeoutError as exc:
|
||||
with contextlib.suppress(ProcessLookupError):
|
||||
process.kill()
|
||||
|
||||
raise TimeoutError(f"Command '{cmd}' timed out after {timeout} seconds") from exc
|
||||
+143
@@ -0,0 +1,143 @@
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import tokenize
|
||||
from pathlib import Path
|
||||
|
||||
from unidiff import PatchSet
|
||||
|
||||
|
||||
def remove_comments_from_line(line: str) -> str:
|
||||
try:
|
||||
tokens = tokenize.generate_tokens(io.StringIO(line).readline)
|
||||
result_parts = []
|
||||
prev_end = (0, 0)
|
||||
|
||||
for tok_type, tok_str, tok_start, tok_end, _ in tokens:
|
||||
if tok_type == tokenize.COMMENT:
|
||||
break
|
||||
(srow, scol) = tok_start
|
||||
if srow == 1 and scol > prev_end[1]:
|
||||
result_parts.append(line[prev_end[1] : scol])
|
||||
result_parts.append(tok_str)
|
||||
prev_end = tok_end
|
||||
|
||||
return "".join(result_parts).rstrip()
|
||||
except tokenize.TokenError:
|
||||
if "#" in line:
|
||||
return line.split("#", 1)[0].rstrip()
|
||||
return line
|
||||
|
||||
|
||||
def clean_patch(ori_patch_text):
|
||||
# in case ori_patch_text has unexpected trailing newline characters
|
||||
# processed_ori_patch_text = ""
|
||||
# previous_line = None
|
||||
# for line in ori_patch_text.split('\n'):
|
||||
# if previous_line is None:
|
||||
# previous_line = line
|
||||
# continue
|
||||
# elif previous_line.strip() == '' and "diff --git" in line:
|
||||
# previous_line = line
|
||||
# continue
|
||||
# else:
|
||||
# processed_ori_patch_text = processed_ori_patch_text + previous_line + "\n"
|
||||
# previous_line = line
|
||||
# if previous_line:
|
||||
# processed_ori_patch_text = processed_ori_patch_text + previous_line
|
||||
|
||||
processed_ori_patch_text = ori_patch_text
|
||||
patch = PatchSet(processed_ori_patch_text)
|
||||
extracted_lines = []
|
||||
delete_lines = []
|
||||
add_lines = []
|
||||
for patched_file in patch:
|
||||
for hunk in patched_file:
|
||||
for line in hunk:
|
||||
if line.is_added:
|
||||
content = line.value.lstrip("+")
|
||||
if content.strip() and not re.match(r"^\s*#", content):
|
||||
content = remove_comments_from_line(content.rstrip())
|
||||
extracted_lines.append("+" + content)
|
||||
add_lines.append(content)
|
||||
elif line.is_removed:
|
||||
content = line.value.lstrip("-")
|
||||
if content.strip() and not re.match(r"^\s*#", content):
|
||||
content = remove_comments_from_line(content.rstrip())
|
||||
extracted_lines.append("-" + content)
|
||||
delete_lines.append(content)
|
||||
new_patch_text = "\n".join(extracted_lines)
|
||||
|
||||
new_patch_text = re.sub(r"\s+", "", new_patch_text)
|
||||
|
||||
return new_patch_text
|
||||
|
||||
|
||||
def save_patches(instance_id, patches_path, patches, group_id=1):
|
||||
trial_index = 1
|
||||
|
||||
dir_path = Path(patches_path) / f"group_{group_id}"
|
||||
dir_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def get_unique_filename(patches_path, trial_index):
|
||||
filename = f"{instance_id}_{trial_index}.patch"
|
||||
while os.path.exists(dir_path / filename):
|
||||
trial_index += 1
|
||||
filename = f"{instance_id}_{trial_index}.patch"
|
||||
return filename
|
||||
|
||||
patch_file = get_unique_filename(patches_path, trial_index)
|
||||
|
||||
clean_patch = patches
|
||||
with open(dir_path / patch_file, "w") as file:
|
||||
file.write(clean_patch)
|
||||
|
||||
print(f"Patches saved in {dir_path / patch_file}")
|
||||
|
||||
|
||||
def get_trajectory_filename(instance_id, traj_dir, group_id=1, voting_id=1):
|
||||
dir_path = Path(traj_dir) / f"group_{group_id}"
|
||||
dir_path.mkdir(parents=True, exist_ok=True)
|
||||
print("dir_path", dir_path)
|
||||
|
||||
def get_unique_filename():
|
||||
trial_index = 1
|
||||
filename = f"{instance_id}_voting_{voting_id}_trail_{trial_index}.json"
|
||||
while os.path.exists(dir_path / filename):
|
||||
trial_index += 1
|
||||
filename = f"{instance_id}_voting_{voting_id}_trail_{trial_index}.json"
|
||||
return filename
|
||||
|
||||
filename = dir_path / get_unique_filename()
|
||||
return filename.absolute().as_posix()
|
||||
|
||||
|
||||
def save_selection_success(
|
||||
instance_id: str,
|
||||
statistics_path: str,
|
||||
patch_id: int,
|
||||
is_success: int,
|
||||
group_id=1,
|
||||
is_all_success=False,
|
||||
is_all_failed=False,
|
||||
):
|
||||
dir_path = Path(statistics_path) / f"group_{group_id}"
|
||||
dir_path.mkdir(parents=True, exist_ok=True)
|
||||
file_path = dir_path / f"{instance_id}.json"
|
||||
|
||||
with open(file_path, "w") as statistics_file:
|
||||
statistics_file.write(
|
||||
json.dumps(
|
||||
{
|
||||
"instance_id": instance_id,
|
||||
"patch_id": patch_id,
|
||||
"is_success": is_success,
|
||||
"is_all_success": is_all_success,
|
||||
"is_all_failed": is_all_failed,
|
||||
},
|
||||
indent=4,
|
||||
sort_keys=True,
|
||||
ensure_ascii=False,
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,474 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import argparse
|
||||
import io
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
import tarfile
|
||||
import traceback
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from docker import DockerClient, from_env
|
||||
from docker.errors import ImageNotFound
|
||||
from docker.models.containers import Container
|
||||
from tqdm import tqdm
|
||||
|
||||
from .utils import BENCHMARK_CONFIG, docker_exec
|
||||
|
||||
|
||||
class BenchmarkEvaluation:
|
||||
"""
|
||||
Main class for running experiments and evaluations.
|
||||
Handles Docker image management, environment preparation, patch generation, and evaluation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
benchmark: str,
|
||||
working_dir: str,
|
||||
trae_config_file_name: str,
|
||||
dataset: str = "SWE-bench_Verified",
|
||||
docker_env_config: str = "",
|
||||
benchmark_harness_path: str = "",
|
||||
run_id: str = "trae-agent",
|
||||
max_workers: int = 4,
|
||||
instance_ids: list[str] | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize the BenchmarkEvaluation class.
|
||||
|
||||
Args:
|
||||
benchmark: Benchmark name.
|
||||
working_dir: Path for workspace (used for temp files and artifacts).
|
||||
trae_config_file_name: Path to Trae config file.
|
||||
dataset: Dataset name.
|
||||
docker_env_config: Path to Docker environment config file.
|
||||
benchmark_harness_path: Path to benchmark harness (for evaluation).
|
||||
run_id: Unique run identifier.
|
||||
max_workers: Maximum number of parallel workers.
|
||||
instance_ids: List of instance IDs to run (optional).
|
||||
"""
|
||||
assert benchmark in BENCHMARK_CONFIG, f"Invalid benchmark name: {benchmark}"
|
||||
self.config = BENCHMARK_CONFIG[benchmark]
|
||||
self.dataset_name = dataset
|
||||
assert self.dataset_name in self.config.valid_datasets, (
|
||||
f"Invalid dataset name: {self.dataset_name}"
|
||||
)
|
||||
|
||||
self.benchmark = benchmark
|
||||
self.dataset = self.config.load_dataset(self.dataset_name)
|
||||
self.docker_client: DockerClient = from_env()
|
||||
self.image_status: dict[Any, Any] = {}
|
||||
|
||||
self.working_dir = Path(working_dir)
|
||||
self.benchmark_harness_path = benchmark_harness_path
|
||||
self.run_id = run_id
|
||||
self.max_workers = max_workers
|
||||
if instance_ids is None:
|
||||
instance_ids = [instance["instance_id"] for instance in self.dataset]
|
||||
else:
|
||||
self.instance_ids = instance_ids
|
||||
|
||||
if docker_env_config != "":
|
||||
with open(docker_env_config, "r") as f:
|
||||
self.docker_env_config: dict[str, dict[str, str]] = json.load(f)
|
||||
else:
|
||||
self.docker_env_config = {}
|
||||
|
||||
self.working_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.trae_config_file_name = trae_config_file_name
|
||||
shutil.copyfile(self.trae_config_file_name, self.working_dir / "trae_config.yaml")
|
||||
|
||||
self.results_dir = Path("results")
|
||||
self.task_id = f"{self.benchmark}_{self.dataset_name}_{self.run_id}".replace("/", "_")
|
||||
self.task_results_dir = self.results_dir / self.task_id
|
||||
self.task_results_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.pull_images()
|
||||
|
||||
def _image_name(self, instance_id: str) -> str:
|
||||
"""
|
||||
Get the Docker image name for a given instance ID.
|
||||
|
||||
Args:
|
||||
instance_id: Instance identifier.
|
||||
|
||||
Returns:
|
||||
Docker image name string.
|
||||
"""
|
||||
return self.config.image_name(instance_id)
|
||||
|
||||
def _check_images(self):
|
||||
"""
|
||||
Check existence of required Docker images for all instances.
|
||||
Updates self.image_status dict.
|
||||
"""
|
||||
for item in tqdm(self.dataset, desc="Checking image status"):
|
||||
instance_id: str = item["instance_id"]
|
||||
image_name = self._image_name(instance_id)
|
||||
try:
|
||||
_ = self.docker_client.images.get(image_name)
|
||||
self.image_status[instance_id] = True
|
||||
except ImageNotFound:
|
||||
self.image_status[instance_id] = False
|
||||
|
||||
try:
|
||||
_ = self.docker_client.images.get("ubuntu:22.04")
|
||||
except Exception:
|
||||
self.docker_client.images.pull("ubuntu:22.04")
|
||||
|
||||
def pull_images(self):
|
||||
"""
|
||||
Pull missing Docker images required for all instances.
|
||||
"""
|
||||
self._check_images()
|
||||
ids = self.instance_ids if self.instance_ids else list(self.image_status.keys())
|
||||
print(f"Total number of images: {len(ids)}")
|
||||
instance_ids = [instance_id for instance_id in ids if not self.image_status[instance_id]]
|
||||
print(f"Number of images to download: {len(instance_ids)}")
|
||||
if len(instance_ids) == 0:
|
||||
return
|
||||
for instance_id in tqdm(instance_ids, desc="Downloading images"):
|
||||
image_name = self._image_name(instance_id)
|
||||
self.docker_client.images.pull(image_name)
|
||||
|
||||
def prepare_trae_agent(self):
|
||||
"""
|
||||
Build Trae Agent and UV inside a base Ubuntu container.
|
||||
Save built artifacts to workspace for later use in experiment containers.
|
||||
"""
|
||||
tars = ["trae-agent.tar", "uv.tar", "uv_shared.tar"]
|
||||
all_exist = all((self.working_dir / tar).exists() for tar in tars)
|
||||
if all_exist:
|
||||
print("Found built trae-agent and uv artifacts. Skipping building.")
|
||||
return
|
||||
|
||||
try:
|
||||
image = self.docker_client.images.get("ubuntu:22.04")
|
||||
except Exception:
|
||||
image = self.docker_client.images.pull("ubuntu:22.04")
|
||||
|
||||
repo_root_path = Path(__file__).parent.parent
|
||||
assert (repo_root_path / "trae_agent" / "__init__.py").is_file()
|
||||
|
||||
container = self.docker_client.containers.run(
|
||||
image=image,
|
||||
command="bash",
|
||||
detach=True,
|
||||
tty=True,
|
||||
stdin_open=True,
|
||||
volumes={
|
||||
self.working_dir.absolute().as_posix(): {"bind": "/trae-workspace", "mode": "rw"},
|
||||
repo_root_path.absolute().as_posix(): {"bind": "/trae-src", "mode": "ro"},
|
||||
},
|
||||
environment=self.docker_env_config.get("preparation_env", None),
|
||||
)
|
||||
|
||||
build_commands = [
|
||||
"apt-get update",
|
||||
"apt-get install -y curl",
|
||||
"curl -LsSf https://astral.sh/uv/install.sh | sh",
|
||||
"rm -rf /trae-workspace/trae-agent && mkdir /trae-workspace/trae-agent",
|
||||
"cp -r -t /trae-workspace/trae-agent/ /trae-src/trae_agent /trae-src/.python-version /trae-src/pyproject.toml /trae-src/uv.lock /trae-src/README.md",
|
||||
"cd /trae-workspace/trae-agent && source $HOME/.local/bin/env && uv sync",
|
||||
]
|
||||
|
||||
for command in tqdm(
|
||||
build_commands, desc="Building trae-agent inside base Docker container"
|
||||
):
|
||||
try:
|
||||
new_command = f'/bin/bash -c "{command}"'
|
||||
return_code, output = docker_exec(container, new_command)
|
||||
except Exception:
|
||||
print(f"{command} failed.")
|
||||
print(traceback.format_exc())
|
||||
break
|
||||
if return_code is not None and return_code != 0:
|
||||
print("Docker exec error. Error message: {}".format(output))
|
||||
container.stop()
|
||||
container.remove()
|
||||
exit(-1)
|
||||
|
||||
for tar_name, src_path in [
|
||||
("trae-agent.tar", "/trae-workspace/trae-agent"),
|
||||
("uv.tar", "/root/.local/bin/uv"),
|
||||
("uv_shared.tar", "/root/.local/share/uv"),
|
||||
]:
|
||||
try:
|
||||
with open(self.working_dir / tar_name, "wb") as f:
|
||||
bits, _ = container.get_archive(src_path)
|
||||
for chunk in bits:
|
||||
f.write(chunk)
|
||||
except Exception:
|
||||
print(f"Failed to save {tar_name} from container.")
|
||||
|
||||
container.stop()
|
||||
container.remove()
|
||||
|
||||
def prepare_experiment_container(self, instance: dict[str, str]) -> Container:
|
||||
"""
|
||||
Prepare experiment Docker container for a given instance.
|
||||
The container mounts the results directory for this instance,
|
||||
so all outputs are directly accessible on the host.
|
||||
Args:
|
||||
instance: Instance dictionary.
|
||||
Returns:
|
||||
Docker container object.
|
||||
"""
|
||||
|
||||
image_name = self._image_name(instance["instance_id"])
|
||||
instance_result_dir = self.task_results_dir / instance["instance_id"]
|
||||
instance_result_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.config.problem_statement(instance, instance_result_dir)
|
||||
|
||||
container: Container = self.docker_client.containers.run(
|
||||
image_name,
|
||||
command="/bin/bash",
|
||||
detach=True,
|
||||
tty=True,
|
||||
stdin_open=True,
|
||||
volumes={
|
||||
instance_result_dir.absolute().as_posix(): {"bind": "/instance-data", "mode": "rw"},
|
||||
},
|
||||
working_dir="/trae-workspace",
|
||||
environment=self.docker_env_config.get("experiment_env", None),
|
||||
stream=True,
|
||||
)
|
||||
|
||||
for fname in ["trae-agent.tar", "uv.tar", "uv_shared.tar", "trae_config.yaml"]:
|
||||
tar_stream = io.BytesIO()
|
||||
with tarfile.open(fileobj=tar_stream, mode="w") as tar:
|
||||
tar.add(self.working_dir / fname, arcname=fname)
|
||||
tar_stream.seek(0)
|
||||
container.put_archive("/trae-workspace", tar_stream.getvalue())
|
||||
|
||||
setup_commands = [
|
||||
"tar xf trae-agent.tar",
|
||||
"tar xf uv.tar",
|
||||
"mkdir -p /root/.local/bin",
|
||||
"mv uv /root/.local/bin/",
|
||||
"tar xf uv_shared.tar",
|
||||
"mkdir -p /root/.local/share",
|
||||
"mv uv /root/.local/share/",
|
||||
]
|
||||
for command in setup_commands:
|
||||
try:
|
||||
new_command = f'/bin/bash -c "{command}"'
|
||||
return_code, output = docker_exec(container, new_command)
|
||||
if return_code is not None and return_code != 0:
|
||||
print("Docker exec error. Error message: {}".format(output))
|
||||
except Exception:
|
||||
print(f"{command} failed.")
|
||||
print(traceback.format_exc())
|
||||
break
|
||||
return container
|
||||
|
||||
def run_one_instance(self, instance_id: str):
|
||||
"""
|
||||
Run patch generation for a single instance.
|
||||
All outputs are written directly to the mounted results directory.
|
||||
Args:
|
||||
instance_id: Instance identifier.
|
||||
"""
|
||||
instance = next((inst for inst in self.dataset if inst["instance_id"] == instance_id), None)
|
||||
if instance is None:
|
||||
print(f"Instance {instance_id} not found.")
|
||||
return
|
||||
|
||||
working_dir = self.config.working_dir(instance_id)
|
||||
|
||||
container_problem_statement_path = "/instance-data/problem_statement.txt"
|
||||
container_patch_file_path = f"/instance-data/{instance_id}.patch"
|
||||
container_traj_path = f"/instance-data/{instance_id}.json"
|
||||
|
||||
container = self.prepare_experiment_container(instance)
|
||||
command = (
|
||||
f"source trae-agent/.venv/bin/activate && "
|
||||
f"trae-cli run --file {container_problem_statement_path} "
|
||||
f'--working-dir="{working_dir}" '
|
||||
f"--config-file trae_config.yaml --must-patch "
|
||||
f"--patch-path {container_patch_file_path} --trajectory-file {container_traj_path}"
|
||||
)
|
||||
new_command = f"/bin/bash -c '{command}'"
|
||||
try:
|
||||
return_code, output = docker_exec(container, new_command)
|
||||
if return_code is not None and return_code != 0:
|
||||
print("Docker exec error. Error message: {}".format(output))
|
||||
except Exception:
|
||||
print(f"{command} failed.")
|
||||
print(traceback.format_exc())
|
||||
|
||||
container.stop()
|
||||
container.remove()
|
||||
|
||||
def run_all(self):
|
||||
"""
|
||||
Run patch generation for all instances in the dataset, with parallelism controlled by max_workers.
|
||||
"""
|
||||
instance_ids = [instance["instance_id"] for instance in self.dataset]
|
||||
with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
|
||||
futures = {
|
||||
executor.submit(self.run_one_instance, instance_id): instance_id
|
||||
for instance_id in instance_ids
|
||||
}
|
||||
for future in tqdm(
|
||||
as_completed(futures), total=len(futures), desc="Running all instances"
|
||||
):
|
||||
instance_id = futures[future]
|
||||
try:
|
||||
future.result()
|
||||
except Exception as e:
|
||||
print(f"Instance {instance_id} failed: {e}")
|
||||
|
||||
def run_eval(self):
|
||||
"""
|
||||
Run evaluation using the benchmark harness.
|
||||
Evaluation results and predictions.json are stored in the task results directory.
|
||||
"""
|
||||
self.config.evaluate_harness_before(
|
||||
self.task_results_dir, self.dataset_name, self.max_workers
|
||||
)
|
||||
|
||||
benchmark_harness_path = Path(self.benchmark_harness_path)
|
||||
cmd = self.config.evaluate_harness(
|
||||
self.dataset_name, self.task_results_dir, self.task_id, self.max_workers
|
||||
)
|
||||
process = subprocess.run(cmd, capture_output=True, cwd=benchmark_harness_path.as_posix())
|
||||
print(process.stdout.decode())
|
||||
print(process.stderr.decode())
|
||||
|
||||
result_filename = "results.json"
|
||||
result_path = self.task_results_dir / result_filename
|
||||
print(f"Evaluation completed and file saved to {result_path}")
|
||||
|
||||
self.config.evaluate_harness_after(self.benchmark_harness_path, self.task_id)
|
||||
|
||||
def get_all_preds(self, instance_ids: list[str] | None = None):
|
||||
"""
|
||||
Collect all generated patches and write predictions.json to results directory.
|
||||
|
||||
Args:
|
||||
instance_ids: List of instance IDs to collect (optional).
|
||||
"""
|
||||
preds: list[dict[str, str]] = []
|
||||
if not instance_ids:
|
||||
instance_ids = [instance["instance_id"] for instance in self.dataset]
|
||||
for instance_id in instance_ids:
|
||||
patch_path = self.task_results_dir / instance_id / f"{instance_id}.patch"
|
||||
if not patch_path.exists():
|
||||
continue
|
||||
with open(patch_path, "r") as f:
|
||||
patch = f.read()
|
||||
preds.append(
|
||||
{
|
||||
"instance_id": instance_id,
|
||||
"model_name_or_path": "trae-agent",
|
||||
"model_patch": patch,
|
||||
}
|
||||
)
|
||||
with open(self.task_results_dir / "predictions.json", "w") as f:
|
||||
json.dump(preds, f)
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
Main entry point for benchmark evaluation script.
|
||||
Parses command-line arguments and runs patch generation and/or evaluation.
|
||||
"""
|
||||
argument_parser = argparse.ArgumentParser()
|
||||
argument_parser.add_argument(
|
||||
"--benchmark", type=str, default="SWE-bench", help="Benchmark name."
|
||||
)
|
||||
argument_parser.add_argument(
|
||||
"--dataset", type=str, default="SWE-bench_Verified", help="Dataset name."
|
||||
)
|
||||
argument_parser.add_argument(
|
||||
"--working-dir", type=str, default="./trae-workspace", help="Workspace directory."
|
||||
)
|
||||
argument_parser.add_argument(
|
||||
"--config-file", type=str, default="trae_config.yaml", help="Trae agent config file path."
|
||||
)
|
||||
argument_parser.add_argument(
|
||||
"--docker-env-config", type=str, default="", required=False, help="Docker env config file."
|
||||
)
|
||||
argument_parser.add_argument(
|
||||
"--instance_ids",
|
||||
nargs="+",
|
||||
type=str,
|
||||
help="Instance IDs to run (space separated).",
|
||||
)
|
||||
argument_parser.add_argument(
|
||||
"--benchmark-harness-path",
|
||||
type=str,
|
||||
default="",
|
||||
required=False,
|
||||
help="Path to benchmark harness (for evaluation).",
|
||||
)
|
||||
argument_parser.add_argument(
|
||||
"--run-id",
|
||||
type=str,
|
||||
required=False,
|
||||
default="trae-agent",
|
||||
help="Run ID for benchmark evaluation.",
|
||||
)
|
||||
argument_parser.add_argument(
|
||||
"--mode",
|
||||
type=str,
|
||||
choices=["e2e", "expr", "eval"],
|
||||
default="e2e",
|
||||
help="e2e: both patch generation and evaluation; expr: only patch generation; eval: only evaluation.",
|
||||
)
|
||||
argument_parser.add_argument(
|
||||
"--max_workers", type=int, default=4, help="Maximum number of parallel workers."
|
||||
)
|
||||
|
||||
args = argument_parser.parse_args()
|
||||
evaluation = BenchmarkEvaluation(
|
||||
args.benchmark,
|
||||
args.working_dir,
|
||||
args.config_file,
|
||||
args.dataset,
|
||||
args.docker_env_config,
|
||||
args.benchmark_harness_path,
|
||||
args.run_id,
|
||||
args.max_workers,
|
||||
args.instance_ids,
|
||||
)
|
||||
|
||||
evaluation.prepare_trae_agent()
|
||||
|
||||
# Patch generation (expr/e2e mode)
|
||||
if args.mode in ("e2e", "expr"):
|
||||
if args.instance_ids:
|
||||
print(f"Running specified instances: {args.instance_ids}")
|
||||
with ThreadPoolExecutor(max_workers=args.max_workers) as executor:
|
||||
futures = {
|
||||
executor.submit(evaluation.run_one_instance, instance_id): instance_id
|
||||
for instance_id in args.instance_ids
|
||||
}
|
||||
for future in tqdm(
|
||||
as_completed(futures), total=len(futures), desc="Running instances"
|
||||
):
|
||||
instance_id = futures[future]
|
||||
try:
|
||||
future.result()
|
||||
except Exception as e:
|
||||
print(f"Instance {instance_id} failed: {e}")
|
||||
else:
|
||||
print("Running all instances in dataset.")
|
||||
evaluation.run_all()
|
||||
|
||||
# Evaluation (eval/e2e mode)
|
||||
if args.mode in ("e2e", "eval"):
|
||||
evaluation.get_all_preds(args.instance_ids)
|
||||
evaluation.run_eval()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+40
@@ -0,0 +1,40 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
set -e
|
||||
|
||||
case "$1" in
|
||||
multi_swe_bench)
|
||||
MULTI_SWE_BENCH_COMMIT_HASH="9a9bec0f3725e1e5340299192571f3a4c26ea27d"
|
||||
git clone https://github.com/multi-swe-bench/multi-swe-bench.git
|
||||
cd multi-swe-bench
|
||||
git checkout $MULTI_SWE_BENCH_COMMIT_HASH
|
||||
python3 -m venv multi_swebench_venv
|
||||
source multi_swebench_venv/bin/activate
|
||||
make install
|
||||
deactivate
|
||||
;;
|
||||
swe_bench)
|
||||
SWE_BENCH_COMMIT_HASH="2bf15e1be3c995a0758529bd29848a8987546090"
|
||||
git clone https://github.com/SWE-bench/SWE-bench.git
|
||||
cd SWE-bench
|
||||
git checkout $SWE_BENCH_COMMIT_HASH
|
||||
python3 -m venv swebench_venv
|
||||
source swebench_venv/bin/activate
|
||||
pip install -e .
|
||||
deactivate
|
||||
;;
|
||||
swe_bench_live)
|
||||
SWE_BENCH_LIVE_COMMIT_HASH="cbc2a3ce1d3d0ce588a45ad6730a04623a84a933"
|
||||
git clone https://github.com/microsoft/SWE-bench-Live.git
|
||||
cd SWE-bench-Live
|
||||
git checkout $SWE_BENCH_LIVE_COMMIT_HASH
|
||||
python3 -m venv swebench_live_venv
|
||||
source swebench_live_venv/bin/activate
|
||||
pip install -e .
|
||||
deactivate
|
||||
;;
|
||||
*)
|
||||
echo "Usage: ./setup.sh [multi_swe_bench|swe_bench|swe_bench_live]"
|
||||
;;
|
||||
esac
|
||||
@@ -0,0 +1,276 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
from datasets import load_dataset
|
||||
from docker.models.containers import Container, ExecResult
|
||||
|
||||
|
||||
def docker_exec(container: Container, command: str):
|
||||
"""
|
||||
Execute a shell command inside a Docker container.
|
||||
|
||||
Args:
|
||||
container: Docker container object.
|
||||
command: Shell command to execute.
|
||||
|
||||
Returns:
|
||||
Tuple (return_code, output_str).
|
||||
"""
|
||||
exec_result: ExecResult = container.exec_run(cmd=command)
|
||||
return_code = exec_result[0]
|
||||
output = exec_result[1].decode("utf-8")
|
||||
return return_code, output
|
||||
|
||||
|
||||
def swebench_evaluate_harness_after(benchmark_harness_path, task_id):
|
||||
src_base = f"{benchmark_harness_path}/logs/run_evaluation/{task_id}/trae-agent"
|
||||
dst_base = f"results/{task_id}"
|
||||
json_src = f"{benchmark_harness_path}/trae-agent.{task_id}.json"
|
||||
json_dst = os.path.join(dst_base, "results.json")
|
||||
if not os.path.exists(src_base):
|
||||
print(f"Source directory does not exist: {src_base}")
|
||||
return
|
||||
for folder_name in os.listdir(src_base):
|
||||
src_folder = os.path.join(src_base, folder_name)
|
||||
dst_folder = os.path.join(dst_base, folder_name)
|
||||
if os.path.isdir(src_folder):
|
||||
os.makedirs(dst_folder, exist_ok=True)
|
||||
for file_name in os.listdir(src_folder):
|
||||
src_file = os.path.join(src_folder, file_name)
|
||||
dst_file = os.path.join(dst_folder, file_name)
|
||||
if not os.path.exists(dst_file):
|
||||
shutil.copy2(src_file, dst_file)
|
||||
os.makedirs(dst_base, exist_ok=True)
|
||||
if not os.path.exists(json_dst):
|
||||
shutil.copy2(json_src, json_dst)
|
||||
|
||||
|
||||
def multi_swebench_evaluate_harness_after(benchmark_harness_path, task_id):
|
||||
task_results_dir = Path("results") / task_id
|
||||
output_dir = (task_results_dir / "dataset").resolve()
|
||||
src_file = output_dir / "final_report.json"
|
||||
dst_file = task_results_dir / "results.json"
|
||||
if not src_file.exists():
|
||||
raise FileNotFoundError(f"{src_file} not found")
|
||||
shutil.copyfile(src_file, dst_file)
|
||||
|
||||
|
||||
def _write_problem_statement(instance_dir: Path, content: str) -> int:
|
||||
"""Helper function to write problem statement using context manager."""
|
||||
with open(instance_dir / "problem_statement.txt", "w", encoding="utf-8") as f:
|
||||
return f.write(content)
|
||||
|
||||
|
||||
def _load_jsonl_dataset(dataset_name: str) -> list[dict]:
|
||||
"""Helper function to load JSONL dataset using context manager."""
|
||||
result = []
|
||||
with open(f"{dataset_name.lower().replace('-', '_')}.jsonl", "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
if line.strip():
|
||||
result.append(json.loads(line))
|
||||
return result
|
||||
|
||||
|
||||
def _write_multi_problem_statement(instance_dir: Path, resolved_issues: list[dict]) -> int:
|
||||
"""Helper function to write multi-issue problem statement using context manager."""
|
||||
content = "\n".join(
|
||||
issue.get("title", "") + "\n" + issue.get("body", "") for issue in resolved_issues
|
||||
)
|
||||
with open(instance_dir / "problem_statement.txt", "w", encoding="utf-8") as f:
|
||||
return f.write(content)
|
||||
|
||||
|
||||
def multi_swebench_evaluate_harness_before(task_results_dir, dataset_name, max_workers):
|
||||
task_results_dir = Path(task_results_dir)
|
||||
pred_json_path = task_results_dir / "predictions.json"
|
||||
pred_jsonl_path = task_results_dir / "predictions.jsonl"
|
||||
dataset_file_path = f"{dataset_name.lower().replace('-', '_')}.jsonl"
|
||||
|
||||
instance_map = {}
|
||||
with open(dataset_file_path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
if not line.strip():
|
||||
continue
|
||||
item = json.loads(line)
|
||||
instance_id = item.get("instance_id")
|
||||
org = item.get("org")
|
||||
repo = item.get("repo")
|
||||
number = item.get("number")
|
||||
instance_map[instance_id] = {"org": org, "repo": repo, "number": number}
|
||||
|
||||
with open(pred_json_path, "r", encoding="utf-8") as f:
|
||||
preds = json.load(f)
|
||||
with open(pred_jsonl_path, "w", encoding="utf-8") as f:
|
||||
for item in preds:
|
||||
instance_id = item["instance_id"]
|
||||
patch = item["model_patch"]
|
||||
info = instance_map.get(instance_id, {})
|
||||
new_item = {
|
||||
"org": info.get("org"),
|
||||
"repo": info.get("repo"),
|
||||
"number": info.get("number"),
|
||||
"fix_patch": patch,
|
||||
}
|
||||
f.write(json.dumps(new_item, ensure_ascii=False) + "\n")
|
||||
|
||||
base_dir = Path(__file__).resolve().parent
|
||||
task_results_dir = base_dir / task_results_dir
|
||||
patch_file_path = str((base_dir / pred_jsonl_path).resolve())
|
||||
dataset_file_path = str((base_dir / dataset_file_path).resolve())
|
||||
|
||||
output_dir = (task_results_dir / "dataset").resolve()
|
||||
repo_dir = (task_results_dir / "repos").resolve()
|
||||
log_dir = (task_results_dir / "logs").resolve()
|
||||
workdir = (task_results_dir / "workdir").resolve()
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
repo_dir.mkdir(parents=True, exist_ok=True)
|
||||
log_dir.mkdir(parents=True, exist_ok=True)
|
||||
workdir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
output_dir = str(output_dir)
|
||||
repo_dir = str(repo_dir)
|
||||
log_dir = str(log_dir)
|
||||
workdir = str(workdir)
|
||||
|
||||
config = {
|
||||
"mode": "evaluation",
|
||||
"workdir": workdir,
|
||||
"patch_files": [patch_file_path],
|
||||
"dataset_files": [dataset_file_path],
|
||||
"force_build": False,
|
||||
"output_dir": output_dir,
|
||||
"specifics": [],
|
||||
"skips": [],
|
||||
"repo_dir": repo_dir,
|
||||
"need_clone": False,
|
||||
"global_env": [],
|
||||
"clear_env": True,
|
||||
"stop_on_error": True,
|
||||
"max_workers": max_workers,
|
||||
"max_workers_build_image": max_workers,
|
||||
"max_workers_run_instance": max_workers,
|
||||
"log_dir": log_dir,
|
||||
"log_level": "DEBUG",
|
||||
}
|
||||
|
||||
config_path = task_results_dir / "evaluate_config.json"
|
||||
with open(config_path, "w", encoding="utf-8") as f:
|
||||
json.dump(config, f, indent=2)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkConfig:
|
||||
valid_datasets: list[str]
|
||||
load_dataset: Callable[[str], Any]
|
||||
image_name: Callable[[str], str]
|
||||
problem_statement: Callable[[dict, Path], Any]
|
||||
working_dir: Callable[[str], str]
|
||||
evaluate_harness: Callable[..., list[str]]
|
||||
evaluate_harness_before: Callable[..., Any]
|
||||
evaluate_harness_after: Callable[..., Any]
|
||||
|
||||
|
||||
BENCHMARK_CONFIG: dict[str, BenchmarkConfig] = {
|
||||
# SWE-bench
|
||||
"SWE-bench": BenchmarkConfig(
|
||||
valid_datasets=["SWE-bench", "SWE-bench_Lite", "SWE-bench_Verified"],
|
||||
load_dataset=lambda dataset_name: load_dataset(
|
||||
f"princeton-nlp/{dataset_name}", split="test"
|
||||
),
|
||||
image_name=lambda instance_id: (
|
||||
f"swebench/sweb.eval.x86_64.{instance_id.lower()}:latest".replace("__", "_1776_")
|
||||
),
|
||||
problem_statement=lambda instance, instance_dir: (
|
||||
_write_problem_statement(instance_dir, instance.get("problem_statement", ""))
|
||||
),
|
||||
working_dir=lambda instance_id: "/testbed/",
|
||||
evaluate_harness=lambda dataset_name, task_results_dir, task_id, max_workers: [
|
||||
"swebench_venv/bin/python",
|
||||
"-m",
|
||||
"swebench.harness.run_evaluation",
|
||||
"--dataset_name",
|
||||
f"princeton-nlp/{dataset_name}",
|
||||
"--predictions_path",
|
||||
(task_results_dir / "predictions.json").absolute().as_posix(),
|
||||
"--max_workers",
|
||||
str(max_workers),
|
||||
"--run_id",
|
||||
task_id,
|
||||
"--cache_level",
|
||||
"instance",
|
||||
"--instance_image_tag",
|
||||
"latest",
|
||||
],
|
||||
evaluate_harness_before=lambda *args, **kwargs: None,
|
||||
evaluate_harness_after=swebench_evaluate_harness_after,
|
||||
),
|
||||
# SWE-bench-Live
|
||||
"SWE-bench-Live": BenchmarkConfig(
|
||||
valid_datasets=["SWE-bench-Live/lite", "SWE-bench-Live/verified", "SWE-bench-Live/full"],
|
||||
load_dataset=lambda dataset_name: load_dataset(
|
||||
"SWE-bench-Live/SWE-bench-Live", split=dataset_name.split("/")[-1]
|
||||
),
|
||||
image_name=lambda instance_id: (
|
||||
f"starryzhang/sweb.eval.x86_64.{instance_id.lower()}:latest".replace("__", "_1776_")
|
||||
),
|
||||
problem_statement=lambda instance, instance_dir: (
|
||||
_write_problem_statement(instance_dir, instance.get("problem_statement", ""))
|
||||
),
|
||||
working_dir=lambda instance_id: "/testbed/",
|
||||
evaluate_harness=lambda dataset_name, task_results_dir, task_id, max_workers: [
|
||||
"swebench_live_venv/bin/python",
|
||||
"-m",
|
||||
"swebench.harness.run_evaluation",
|
||||
"--dataset_name",
|
||||
"SWE-bench-Live/SWE-bench-Live",
|
||||
"--namespace",
|
||||
"starryzhang",
|
||||
"--split",
|
||||
dataset_name.split("/")[-1],
|
||||
"--predictions_path",
|
||||
(task_results_dir / "predictions.json").absolute().as_posix(),
|
||||
"--run_id",
|
||||
task_id,
|
||||
"--max_workers",
|
||||
str(max_workers),
|
||||
],
|
||||
evaluate_harness_before=lambda *args, **kwargs: None,
|
||||
evaluate_harness_after=swebench_evaluate_harness_after,
|
||||
),
|
||||
# Multi-SWE-bench
|
||||
"Multi-SWE-bench": BenchmarkConfig(
|
||||
valid_datasets=["Multi-SWE-bench-flash", "Multi-SWE-bench_mini"],
|
||||
load_dataset=lambda dataset_name: _load_jsonl_dataset(dataset_name),
|
||||
image_name=lambda instance_id: (
|
||||
(lambda key: key.rpartition("-")[0] + ":pr-" + key.rpartition("-")[2])(
|
||||
f"mswebench/{instance_id.lower()}".replace("__", "_m_")
|
||||
)
|
||||
),
|
||||
problem_statement=lambda instance, instance_dir: (
|
||||
_write_multi_problem_statement(instance_dir, instance.get("resolved_issues", []))
|
||||
),
|
||||
working_dir=lambda instance_id: (
|
||||
f"/home/{'-'.join(instance_id.split('__')[-1].split('-')[:-1])}/"
|
||||
),
|
||||
evaluate_harness=lambda dataset_name, task_results_dir, task_id, max_workers: [
|
||||
"multi_swebench_venv/bin/python",
|
||||
"-m",
|
||||
"multi_swe_bench.harness.run_evaluation",
|
||||
"--config",
|
||||
os.path.join(
|
||||
os.path.dirname(os.path.abspath(__file__)),
|
||||
task_results_dir / "evaluate_config.json",
|
||||
),
|
||||
],
|
||||
evaluate_harness_before=multi_swebench_evaluate_harness_before,
|
||||
evaluate_harness_after=multi_swebench_evaluate_harness_after,
|
||||
),
|
||||
}
|
||||
+101
@@ -0,0 +1,101 @@
|
||||
[project]
|
||||
name = "trae-agent"
|
||||
version = "0.1.0"
|
||||
description = "LLM-based agent for general purpose software engineering tasks"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"openai>=1.86.0",
|
||||
"anthropic>=0.54.0,<=0.60.0",
|
||||
"click>=8.0.0",
|
||||
"google-genai>=1.24.0",
|
||||
"jsonpath-ng>=1.7.0",
|
||||
"pydantic>=2.0.0",
|
||||
"python-dotenv>=1.0.0",
|
||||
"rich>=13.0.0",
|
||||
"typing-extensions>=4.0.0",
|
||||
"ollama>=0.5.1",
|
||||
"socksio>=1.0.0",
|
||||
"tree-sitter-languages==1.10.2",
|
||||
"tree-sitter==0.21.3",
|
||||
"ruff>=0.12.4",
|
||||
"mcp==1.12.2",
|
||||
"asyncclick>=8.0.0",
|
||||
"pyyaml>=6.0.2",
|
||||
"textual>=0.50.0",
|
||||
"pyinstaller==6.15.0"
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
test = [
|
||||
"pytest>=8.0.0",
|
||||
"pytest-asyncio>=0.23.0",
|
||||
"pytest-mock>=3.12.0",
|
||||
"pytest-cov>=4.0.0",
|
||||
"pre-commit>=4.2.0",
|
||||
]
|
||||
evaluation = [
|
||||
"datasets>=3.6.0",
|
||||
"docker>=7.1.0",
|
||||
"pexpect>=4.9.0",
|
||||
"unidiff>=0.7.5",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
trae-cli = "trae_agent.cli:main"
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["trae_agent"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
minversion = "6.0"
|
||||
addopts = "-ra -q --strict-markers"
|
||||
testpaths = [
|
||||
"tests",
|
||||
]
|
||||
asyncio_mode = "auto"
|
||||
markers = [
|
||||
"slow: marks tests as slow (deselect with '-m \"not slow\"')",
|
||||
"integration: marks tests as integration tests",
|
||||
"unit: marks tests as unit tests",
|
||||
]
|
||||
|
||||
[tool.coverage.run]
|
||||
source = ["trae_agent"]
|
||||
omit = ["tests/*"]
|
||||
|
||||
[tool.coverage.report]
|
||||
exclude_lines = [
|
||||
"pragma: no cover",
|
||||
"def __repr__",
|
||||
"if self.debug:",
|
||||
"if settings.DEBUG",
|
||||
"raise AssertionError",
|
||||
"raise NotImplementedError",
|
||||
"if 0:",
|
||||
"if __name__ == .__main__.:",
|
||||
"class .*\\bProtocol\\):",
|
||||
"@(abc\\.)?abstractmethod",
|
||||
]
|
||||
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 100
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = [
|
||||
"B",
|
||||
"SIM",
|
||||
"C4",
|
||||
"E4", "E9", "E7", "F",
|
||||
"I"
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"types-pyyaml>=6.0.12.20250516",
|
||||
]
|
||||
@@ -0,0 +1,20 @@
|
||||
# HTTP Server
|
||||
|
||||
This folder contains the elements for hosting the Trae agent as an HTTP server using FastAPI. It is still under construction and should **not** be used in production yet.
|
||||
|
||||
## Expected Features of the HTTP Server
|
||||
|
||||
1. The server should be able to perform stateless operations.
|
||||
2. The server should be able to handle concurrent requests.
|
||||
3. The server should always respond in JSON format, even if the response is streaming.
|
||||
|
||||
## Additional Features Expected
|
||||
|
||||
1. The server should be able to reproduce or repeat actions based on a specific JSON file. For example, given a trajectory, it could reproduce specific steps and follow new steps produced by another model.
|
||||
2. To ensure requests are dynamic, the server should support different models, different requests, and different output formats based on the request JSON file.
|
||||
|
||||
## Roadmap
|
||||
|
||||
1. To build a fully functional HTTP Trae agent, we need to gradually split the `trae_agent` into more component-based modules and add more features to make it more dynamic. A specific task is to see if the `run` function can accept an additional parameter called `model`.
|
||||
2. Besides the `run` function, other functions should be callable not only via CLI but also through the HTTP server to meet the second requirement.
|
||||
3. To handle concurrent requests, it is also necessary to ensure the HTTP server is stateless — at least when handling the `run` function operating in different folders.
|
||||
@@ -0,0 +1,151 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from trae_agent.agent.agent_basics import AgentError
|
||||
from trae_agent.agent.trae_agent import TraeAgent
|
||||
from trae_agent.utils.config import Config
|
||||
from trae_agent.utils.legacy_config import LegacyConfig
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMResponse
|
||||
|
||||
|
||||
class TestTraeAgentExtended(unittest.TestCase):
|
||||
def setUp(self):
|
||||
test_config = {
|
||||
"default_provider": "anthropic",
|
||||
"max_steps": 20,
|
||||
"model_providers": {
|
||||
"anthropic": {
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"api_key": "test-dummy-api-key", # dummy api key
|
||||
"max_tokens": 4096,
|
||||
"temperature": 0.5,
|
||||
"top_p": 1,
|
||||
"top_k": 0,
|
||||
"parallel_tool_calls": False,
|
||||
"max_retries": 10,
|
||||
}
|
||||
},
|
||||
}
|
||||
self.config = Config.create_from_legacy_config(legacy_config=LegacyConfig(test_config))
|
||||
|
||||
# Avoid create real LLMClient instance to avoid actual API calls
|
||||
self.llm_client_patcher = patch("trae_agent.agent.base_agent.LLMClient")
|
||||
mock_llm_client = self.llm_client_patcher.start()
|
||||
mock_llm_client.return_value.client = MagicMock()
|
||||
|
||||
if self.config.trae_agent:
|
||||
self.agent = TraeAgent(self.config.trae_agent)
|
||||
else:
|
||||
self.fail("trae_agent config is None")
|
||||
self.test_project_path = "/test/project"
|
||||
self.test_patch_path = "/test/patch.diff"
|
||||
|
||||
def tearDown(self):
|
||||
self.llm_client_patcher.stop()
|
||||
|
||||
def test_new_task_initialization(self):
|
||||
with self.assertRaises(AgentError):
|
||||
self.agent.new_task("test", {}) # Missing required params
|
||||
|
||||
valid_args = {
|
||||
"project_path": self.test_project_path,
|
||||
"issue": "Test issue",
|
||||
"base_commit": "abc123",
|
||||
"must_patch": "true",
|
||||
"patch_path": self.test_patch_path,
|
||||
}
|
||||
self.agent.new_task("test-task", valid_args)
|
||||
|
||||
self.assertEqual(self.agent.project_path, self.test_project_path)
|
||||
self.assertEqual(self.agent.must_patch, "true")
|
||||
self.assertEqual(len(self.agent.tools), 4)
|
||||
self.assertTrue(any(tool.get_name() == "bash" for tool in self.agent.tools))
|
||||
|
||||
@patch("subprocess.check_output")
|
||||
@patch("os.chdir")
|
||||
@patch("os.path.isdir", return_value=True)
|
||||
def test_git_diff_generation(self, mock_isdir, mock_chdir, mock_subprocess):
|
||||
mock_subprocess.return_value = b"test diff"
|
||||
self.agent.project_path = self.test_project_path
|
||||
|
||||
diff = self.agent.get_git_diff()
|
||||
self.assertEqual(diff, "test diff")
|
||||
mock_subprocess.assert_called_with(["git", "--no-pager", "diff"])
|
||||
|
||||
def test_patch_filtering(self):
|
||||
test_patch = """diff --git a/tests/test_example.py b/tests/test_example.py
|
||||
--- a/tests/test_example.py
|
||||
+++ b/tests/test_example.py
|
||||
@@ -5,6 +5,7 @@
|
||||
def test_example(self):
|
||||
assert True
|
||||
"""
|
||||
filtered = self.agent.remove_patches_to_tests(test_patch)
|
||||
self.assertEqual(filtered, "")
|
||||
|
||||
def test_task_completion_detection(self):
|
||||
mock_response = MagicMock(spec=LLMResponse)
|
||||
|
||||
# Test empty patch scenario
|
||||
self.agent.must_patch = "true"
|
||||
self.assertFalse(self.agent._is_task_completed(mock_response))
|
||||
|
||||
# Test valid patch scenario
|
||||
with patch.object(self.agent, "get_git_diff", return_value="valid patch"):
|
||||
self.assertTrue(self.agent._is_task_completed(mock_response))
|
||||
|
||||
def test_tool_initialization(self):
|
||||
tools = [
|
||||
"bash",
|
||||
"str_replace_based_edit_tool",
|
||||
"sequentialthinking",
|
||||
"task_done",
|
||||
]
|
||||
self.agent.new_task("test", {"project_path": self.test_project_path}, tools)
|
||||
tool_names = [tool.get_name() for tool in self.agent.tools]
|
||||
|
||||
self.assertEqual(len(self.agent.tools), len(tools))
|
||||
self.assertIn("bash", tool_names)
|
||||
self.assertIn("str_replace_based_edit_tool", tool_names)
|
||||
self.assertIn("sequentialthinking", tool_names)
|
||||
self.assertIn("task_done", tool_names)
|
||||
|
||||
def test_protected_attributes_access_restrictions(self):
|
||||
"""Test that protected attributes cannot be accessed directly from outside the class."""
|
||||
|
||||
# Test that accessing protected attributes raises AttributeError
|
||||
with self.assertRaises(AttributeError):
|
||||
self.agent.llm_client = 5
|
||||
|
||||
with self.assertRaises(AttributeError):
|
||||
self.agent.max_steps = None
|
||||
|
||||
with self.assertRaises(AttributeError):
|
||||
self.agent.model_config = False
|
||||
|
||||
with self.assertRaises(AttributeError):
|
||||
self.agent.initial_messages = "random"
|
||||
|
||||
with self.assertRaises(AttributeError):
|
||||
_ = self.agent.tool_caller
|
||||
|
||||
def test_public_property_access_allowed(self):
|
||||
"""Test that public properties can be accessed properly."""
|
||||
|
||||
# Test that public properties work correctly
|
||||
self.assertIsNotNone(self.agent.llm_client)
|
||||
self.assertIsNone(self.agent.cli_console)
|
||||
|
||||
# Test that public property setters work
|
||||
from trae_agent.utils.cli import CLIConsole
|
||||
|
||||
mock_console = MagicMock(spec=CLIConsole)
|
||||
self.agent.set_cli_console(mock_console)
|
||||
self.assertEqual(self.agent.cli_console, mock_console)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,181 @@
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from click.testing import CliRunner
|
||||
|
||||
from trae_agent.cli import cli
|
||||
|
||||
|
||||
class TestCli(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.runner = CliRunner()
|
||||
|
||||
@patch("trae_agent.cli.resolve_config_file", return_value="test_config.yaml")
|
||||
@patch("trae_agent.cli.Agent")
|
||||
@patch("trae_agent.cli.asyncio.run")
|
||||
@patch("trae_agent.cli.Config.create")
|
||||
@patch("trae_agent.cli.ConsoleFactory.create_console")
|
||||
def test_run_with_long_prompt(
|
||||
self,
|
||||
mock_create_console,
|
||||
mock_config_create,
|
||||
mock_asyncio_run,
|
||||
mock_agent_class,
|
||||
mock_resolve_config_file,
|
||||
):
|
||||
"""Test that a long prompt string is handled correctly."""
|
||||
# Setup mocks
|
||||
mock_config = MagicMock()
|
||||
mock_config.trae_agent = MagicMock()
|
||||
mock_config_create.return_value.resolve_config_values.return_value = mock_config
|
||||
mock_agent = MagicMock()
|
||||
mock_agent_class.return_value = mock_agent
|
||||
mock_console = MagicMock()
|
||||
# Add the methods that hasattr checks for
|
||||
mock_console.set_initial_task = MagicMock()
|
||||
mock_console.set_agent_context = MagicMock()
|
||||
mock_create_console.return_value = mock_console
|
||||
|
||||
long_prompt = "a" * 500 # A string longer than typical filename limits
|
||||
result = self.runner.invoke(cli, ["run", long_prompt, "--working-dir", "/tmp"])
|
||||
self.assertEqual(result.exit_code, 0)
|
||||
|
||||
# Verify agent.run was called with the long prompt
|
||||
mock_asyncio_run.assert_called_once()
|
||||
mock_agent.run.assert_called_once()
|
||||
args, _ = mock_agent.run.call_args
|
||||
self.assertEqual(args[0], long_prompt)
|
||||
|
||||
@patch("trae_agent.cli.resolve_config_file", return_value="test_config.yaml")
|
||||
@patch("trae_agent.cli.Agent")
|
||||
@patch("trae_agent.cli.asyncio.run")
|
||||
@patch("trae_agent.cli.Config.create")
|
||||
@patch("trae_agent.cli.ConsoleFactory.create_console")
|
||||
def test_run_with_file_argument(
|
||||
self,
|
||||
mock_create_console,
|
||||
mock_config_create,
|
||||
mock_asyncio_run,
|
||||
mock_agent_class,
|
||||
mock_resolve_config_file,
|
||||
):
|
||||
"""Test that the --file argument correctly reads from a file."""
|
||||
# Setup mocks
|
||||
mock_config = MagicMock()
|
||||
mock_config.trae_agent = MagicMock()
|
||||
mock_config_create.return_value.resolve_config_values.return_value = mock_config
|
||||
mock_agent = MagicMock()
|
||||
mock_agent_class.return_value = mock_agent
|
||||
mock_console = MagicMock()
|
||||
# Add the methods that hasattr checks for
|
||||
mock_console.set_initial_task = MagicMock()
|
||||
mock_console.set_agent_context = MagicMock()
|
||||
mock_create_console.return_value = mock_console
|
||||
|
||||
with self.runner.isolated_filesystem():
|
||||
with open("task.txt", "w") as f:
|
||||
f.write("task from file")
|
||||
|
||||
result = self.runner.invoke(cli, ["run", "--file", "task.txt", "--working-dir", "/tmp"])
|
||||
self.assertEqual(result.exit_code, 0)
|
||||
|
||||
# Verify agent.run was called with the file content
|
||||
mock_asyncio_run.assert_called_once()
|
||||
mock_agent.run.assert_called_once()
|
||||
args, _ = mock_agent.run.call_args
|
||||
self.assertEqual(args[0], "task from file")
|
||||
|
||||
@patch("trae_agent.cli.resolve_config_file", return_value="test_config.yaml")
|
||||
def test_run_with_nonexistent_file(self, mock_resolve_config_file):
|
||||
"""Test for a clear error when --file points to a non-existent file."""
|
||||
result = self.runner.invoke(cli, ["run", "--file", "nonexistent.txt"])
|
||||
self.assertNotEqual(result.exit_code, 0)
|
||||
self.assertIn("Error: File not found: nonexistent.txt", result.output)
|
||||
|
||||
@patch("trae_agent.cli.resolve_config_file", return_value="test_config.yaml")
|
||||
def test_run_with_both_task_and_file(self, mock_resolve_config_file):
|
||||
"""Test for a clear error when both task string and --file are used."""
|
||||
result = self.runner.invoke(cli, ["run", "some task", "--file", "task.txt"])
|
||||
self.assertNotEqual(result.exit_code, 0)
|
||||
self.assertIn(
|
||||
"Error: Cannot use both a task string and the --file argument.", result.output
|
||||
)
|
||||
|
||||
def test_run_with_no_input(self):
|
||||
"""Test for a clear error when neither task string nor --file is provided."""
|
||||
result = self.runner.invoke(cli, ["run"])
|
||||
self.assertIn("Error: Config file not found.", result.output)
|
||||
|
||||
@patch("trae_agent.cli.resolve_config_file", return_value="test_config.yaml")
|
||||
@patch("trae_agent.cli.Agent")
|
||||
@patch("trae_agent.cli.Config.create")
|
||||
@patch("trae_agent.cli.ConsoleFactory.create_console")
|
||||
@patch("trae_agent.cli.os.chdir", side_effect=FileNotFoundError("No such file or directory"))
|
||||
def test_run_with_nonexistent_working_dir(
|
||||
self,
|
||||
mock_chdir,
|
||||
mock_create_console,
|
||||
mock_config_create,
|
||||
mock_agent_class,
|
||||
mock_resolve_config_file,
|
||||
):
|
||||
"""Test for a clear error when --working-dir points to a non-existent directory."""
|
||||
# Setup mocks
|
||||
mock_config = MagicMock()
|
||||
mock_config.trae_agent = MagicMock()
|
||||
mock_config_create.return_value.resolve_config_values.return_value = mock_config
|
||||
mock_agent = MagicMock()
|
||||
mock_agent_class.return_value = mock_agent
|
||||
mock_console = MagicMock()
|
||||
mock_console.set_initial_task = MagicMock()
|
||||
mock_console.set_agent_context = MagicMock()
|
||||
mock_create_console.return_value = mock_console
|
||||
|
||||
result = self.runner.invoke(
|
||||
cli, ["run", "some task", "--working-dir", "/path/to/nonexistent/dir"]
|
||||
)
|
||||
self.assertNotEqual(result.exit_code, 0)
|
||||
self.assertIn("Error changing directory", result.output)
|
||||
|
||||
@patch("trae_agent.cli.resolve_config_file", return_value="test_config.yaml")
|
||||
@patch("trae_agent.cli.Agent")
|
||||
@patch("trae_agent.cli.asyncio.run")
|
||||
@patch("trae_agent.cli.Config.create")
|
||||
@patch("trae_agent.cli.ConsoleFactory.create_console")
|
||||
def test_run_with_string_that_is_also_a_filename(
|
||||
self,
|
||||
mock_create_console,
|
||||
mock_config_create,
|
||||
mock_asyncio_run,
|
||||
mock_agent_class,
|
||||
mock_resolve_config_file,
|
||||
):
|
||||
"""Test that a task string that looks like a file is treated as a string."""
|
||||
# Setup mocks
|
||||
mock_config = MagicMock()
|
||||
mock_config.trae_agent = MagicMock()
|
||||
mock_config_create.return_value.resolve_config_values.return_value = mock_config
|
||||
mock_agent = MagicMock()
|
||||
mock_agent_class.return_value = mock_agent
|
||||
mock_console = MagicMock()
|
||||
# Add the methods that hasattr checks for
|
||||
mock_console.set_initial_task = MagicMock()
|
||||
mock_console.set_agent_context = MagicMock()
|
||||
mock_create_console.return_value = mock_console
|
||||
|
||||
with self.runner.isolated_filesystem():
|
||||
with open("task.txt", "w") as f:
|
||||
f.write("file content")
|
||||
|
||||
result = self.runner.invoke(cli, ["run", "task.txt", "--working-dir", "/tmp"])
|
||||
self.assertEqual(result.exit_code, 0)
|
||||
|
||||
# Verify agent.run was called with the string "task.txt", not the file content
|
||||
mock_asyncio_run.assert_called_once()
|
||||
mock_agent.run.assert_called_once()
|
||||
args, _ = mock_agent.run.call_args
|
||||
self.assertEqual(args[0], "task.txt")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,68 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import unittest
|
||||
|
||||
from trae_agent.tools.base import ToolCallArguments
|
||||
from trae_agent.tools.bash_tool import BashTool
|
||||
|
||||
|
||||
class TestBashTool(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.tool = BashTool()
|
||||
|
||||
async def asyncTearDown(self):
|
||||
# Cleanup any active session
|
||||
if self.tool._session:
|
||||
await self.tool._session.stop()
|
||||
|
||||
async def test_tool_initialization(self):
|
||||
self.assertEqual(self.tool.get_name(), "bash")
|
||||
self.assertIn("Run commands in a bash shell", self.tool.get_description())
|
||||
|
||||
params = self.tool.get_parameters()
|
||||
param_names = [p.name for p in params]
|
||||
self.assertIn("command", param_names)
|
||||
self.assertIn("restart", param_names)
|
||||
|
||||
async def test_command_error_handling(self):
|
||||
result = await self.tool.execute(ToolCallArguments({"command": "invalid_command_123"}))
|
||||
|
||||
# Fix assertion: Check if error message contains 'not found' or 'not recognized' (Windows system)
|
||||
self.assertTrue(any(s in result.error.lower() for s in ["not found", "not recognized"]))
|
||||
self.assertNotEqual(result.error_code, 0)
|
||||
|
||||
async def test_session_restart(self):
|
||||
# Ensure session is initialized
|
||||
await self.tool.execute(ToolCallArguments({"command": "echo first session"}))
|
||||
|
||||
# Fix: Check if session object exists
|
||||
self.assertIsNotNone(self.tool._session)
|
||||
|
||||
# Restart and test new session
|
||||
restart_result = await self.tool.execute(ToolCallArguments({"restart": True}))
|
||||
self.assertIn("restarted", restart_result.output.lower())
|
||||
|
||||
# Fix: Ensure new session is created
|
||||
self.assertIsNotNone(self.tool._session)
|
||||
|
||||
# Verify new session works
|
||||
result = await self.tool.execute(ToolCallArguments({"command": "echo new session"}))
|
||||
self.assertIn("new session", result.output)
|
||||
|
||||
async def test_successful_command_execution(self):
|
||||
result = await self.tool.execute(ToolCallArguments({"command": "echo hello world"}))
|
||||
|
||||
# Fix: Check if return code is 0
|
||||
self.assertEqual(result.error_code, 0)
|
||||
self.assertIn("hello world", result.output)
|
||||
self.assertEqual(result.error, "")
|
||||
|
||||
async def test_missing_command_handling(self):
|
||||
result = await self.tool.execute(ToolCallArguments({}))
|
||||
self.assertIn("no command provided", result.error.lower())
|
||||
self.assertEqual(result.error_code, -1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,131 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from trae_agent.tools.base import ToolCallArguments
|
||||
from trae_agent.tools.edit_tool import TextEditorTool
|
||||
|
||||
|
||||
class TestTextEditorTool(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.tool = TextEditorTool()
|
||||
# Use current working directory for test paths
|
||||
self.test_dir = Path.cwd() / "test_dir"
|
||||
self.test_file = self.test_dir / "test_file.txt"
|
||||
|
||||
def mock_file_system(self, exists=True, is_dir=False, content=""):
|
||||
"""Helper to mock file system operations"""
|
||||
patcher = patch("pathlib.Path.exists", return_value=exists)
|
||||
self.mock_exists = patcher.start()
|
||||
self.addCleanup(patcher.stop)
|
||||
|
||||
patcher = patch("pathlib.Path.is_dir", return_value=is_dir)
|
||||
self.mock_is_dir = patcher.start()
|
||||
self.addCleanup(patcher.stop)
|
||||
|
||||
patcher = patch("pathlib.Path.read_text", return_value=content)
|
||||
self.mock_read = patcher.start()
|
||||
self.addCleanup(patcher.stop)
|
||||
|
||||
patcher = patch("pathlib.Path.write_text")
|
||||
self.mock_write = patcher.start()
|
||||
self.addCleanup(patcher.stop)
|
||||
|
||||
async def test_create_file(self):
|
||||
self.mock_file_system(exists=False)
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments(
|
||||
{
|
||||
"command": "create",
|
||||
"path": str(self.test_file),
|
||||
"file_text": "new content",
|
||||
}
|
||||
)
|
||||
)
|
||||
self.mock_write.assert_called_once_with("new content")
|
||||
self.assertIn("created successfully", result.output)
|
||||
|
||||
async def test_insert_line(self):
|
||||
self.mock_file_system(content="line1\nline3")
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments(
|
||||
{
|
||||
"command": "insert",
|
||||
"path": str(self.test_file),
|
||||
"insert_line": 1,
|
||||
"new_str": "line2",
|
||||
}
|
||||
)
|
||||
)
|
||||
self.mock_write.assert_called_once()
|
||||
self.assertIn("edited", result.output)
|
||||
|
||||
async def test_invalid_command(self):
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments({"command": "invalid", "path": str(self.test_file.absolute())})
|
||||
)
|
||||
self.assertEqual(result.error_code, -1)
|
||||
self.assertIn("Please provide a valid path", result.error)
|
||||
|
||||
async def test_str_replace_multiple_occurrences(self):
|
||||
self.mock_file_system(content="dup\ndup\nline3")
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments(
|
||||
{
|
||||
"command": "str_replace",
|
||||
"path": str(self.test_file),
|
||||
"old_str": "dup",
|
||||
"new_str": "new",
|
||||
}
|
||||
)
|
||||
)
|
||||
self.assertEqual(result.error_code, -1)
|
||||
self.assertIn("Multiple occurrences", result.error or "")
|
||||
|
||||
async def test_str_replace_success(self):
|
||||
self.mock_file_system(content="old_content\nline2")
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments(
|
||||
{
|
||||
"command": "str_replace",
|
||||
"path": str(self.test_file),
|
||||
"old_str": "old_content",
|
||||
"new_str": "new_content",
|
||||
}
|
||||
)
|
||||
)
|
||||
self.mock_write.assert_called_once()
|
||||
self.assertIn("edited", result.output)
|
||||
|
||||
async def test_view_directory(self):
|
||||
self.mock_file_system(exists=True, is_dir=True)
|
||||
with patch("trae_agent.tools.edit_tool.run", new_callable=AsyncMock) as mock_run:
|
||||
mock_run.return_value = (0, "file1\nfile2", "")
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments({"command": "view", "path": str(self.test_dir)})
|
||||
)
|
||||
self.assertIn("files and directories", result.output)
|
||||
|
||||
async def test_view_file(self):
|
||||
self.mock_file_system(exists=True, is_dir=False, content="line1\nline2\nline3")
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments({"command": "view", "path": str(self.test_file)})
|
||||
)
|
||||
self.assertRegex(result.output, r"\d+\s+line1")
|
||||
|
||||
async def test_relative_path(self):
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments({"command": "view", "path": "relative/path"})
|
||||
)
|
||||
self.assertIn("absolute path", result.error)
|
||||
|
||||
async def test_missing_parameters(self):
|
||||
result = await self.tool.execute(ToolCallArguments({"command": "create"}))
|
||||
self.assertIn("No path provided", result.error)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,204 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Tests for JSONEditTool."""
|
||||
|
||||
import json
|
||||
import unittest
|
||||
from unittest.mock import mock_open, patch
|
||||
|
||||
from trae_agent.tools.base import ToolCallArguments
|
||||
from trae_agent.tools.json_edit_tool import JSONEditTool
|
||||
|
||||
|
||||
class TestJSONEditTool(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
"""Set up the test environment."""
|
||||
self.tool = JSONEditTool()
|
||||
self.test_file_path = "/test_dir/test_file.json"
|
||||
|
||||
# Default sample data
|
||||
self.sample_data = {
|
||||
"users": [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}],
|
||||
"config": {"enabled": True},
|
||||
}
|
||||
|
||||
def mock_file_read(self, json_data=None):
|
||||
"""Helper to mock file reading operations."""
|
||||
if json_data is None:
|
||||
json_data = self.sample_data
|
||||
|
||||
read_content = json.dumps(json_data)
|
||||
m_open = mock_open(read_data=read_content)
|
||||
|
||||
# Patch open and path checks
|
||||
self.open_patcher = patch("builtins.open", m_open)
|
||||
self.exists_patcher = patch("pathlib.Path.exists", return_value=True)
|
||||
self.is_absolute_patcher = patch("pathlib.Path.is_absolute", return_value=True)
|
||||
|
||||
self.open_patcher.start()
|
||||
self.exists_patcher.start()
|
||||
self.is_absolute_patcher.start()
|
||||
|
||||
self.addCleanup(self.open_patcher.stop)
|
||||
self.addCleanup(self.exists_patcher.stop)
|
||||
self.addCleanup(self.is_absolute_patcher.stop)
|
||||
|
||||
@patch("json.dump")
|
||||
async def test_set_config_value(self, mock_json_dump):
|
||||
"""Test setting a simple configuration value."""
|
||||
self.mock_file_read()
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments(
|
||||
{
|
||||
"operation": "set",
|
||||
"file_path": self.test_file_path,
|
||||
"json_path": "$.config.enabled",
|
||||
"value": False,
|
||||
}
|
||||
)
|
||||
)
|
||||
self.assertEqual(result.error_code, 0)
|
||||
|
||||
# Verify that json.dump was called with the correct data
|
||||
mock_json_dump.assert_called_once()
|
||||
written_data = mock_json_dump.call_args[0][0]
|
||||
self.assertFalse(written_data["config"]["enabled"])
|
||||
|
||||
@patch("json.dump")
|
||||
async def test_update_user_name(self, mock_json_dump):
|
||||
"""Test updating a name in a list of objects."""
|
||||
self.mock_file_read()
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments(
|
||||
{
|
||||
"operation": "set",
|
||||
"file_path": self.test_file_path,
|
||||
"json_path": "$.users[0].name",
|
||||
"value": "Alicia",
|
||||
}
|
||||
)
|
||||
)
|
||||
self.assertEqual(result.error_code, 0)
|
||||
|
||||
mock_json_dump.assert_called_once()
|
||||
written_data = mock_json_dump.call_args[0][0]
|
||||
self.assertEqual(written_data["users"][0]["name"], "Alicia")
|
||||
|
||||
@patch("json.dump")
|
||||
async def test_add_new_user(self, mock_json_dump):
|
||||
"""Test adding a new object to a list (by inserting at the end)."""
|
||||
self.mock_file_read()
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments(
|
||||
{
|
||||
"operation": "add",
|
||||
"file_path": self.test_file_path,
|
||||
"json_path": "$.users[2]", # Inserting at index 2 (end of list)
|
||||
"value": {"id": 3, "name": "Charlie"},
|
||||
}
|
||||
)
|
||||
)
|
||||
self.assertEqual(result.error_code, 0)
|
||||
|
||||
mock_json_dump.assert_called_once()
|
||||
written_data = mock_json_dump.call_args[0][0]
|
||||
self.assertEqual(len(written_data["users"]), 3)
|
||||
self.assertEqual(written_data["users"][2]["name"], "Charlie")
|
||||
|
||||
@patch("json.dump")
|
||||
async def test_add_new_config_key(self, mock_json_dump):
|
||||
"""Test adding a new key-value pair to an object."""
|
||||
self.mock_file_read()
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments(
|
||||
{
|
||||
"operation": "add",
|
||||
"file_path": self.test_file_path,
|
||||
"json_path": "$.config.version",
|
||||
"value": "1.1.0",
|
||||
}
|
||||
)
|
||||
)
|
||||
self.assertEqual(result.error_code, 0)
|
||||
|
||||
mock_json_dump.assert_called_once()
|
||||
written_data = mock_json_dump.call_args[0][0]
|
||||
self.assertEqual(written_data["config"]["version"], "1.1.0")
|
||||
|
||||
@patch("json.dump")
|
||||
async def test_remove_user_by_index(self, mock_json_dump):
|
||||
"""Test removing an element from a list by its index."""
|
||||
self.mock_file_read()
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments(
|
||||
{
|
||||
"operation": "remove",
|
||||
"file_path": self.test_file_path,
|
||||
"json_path": "$.users[0]",
|
||||
}
|
||||
)
|
||||
)
|
||||
self.assertEqual(result.error_code, 0)
|
||||
|
||||
mock_json_dump.assert_called_once()
|
||||
written_data = mock_json_dump.call_args[0][0]
|
||||
self.assertEqual(len(written_data["users"]), 1)
|
||||
self.assertEqual(written_data["users"][0]["name"], "Bob")
|
||||
|
||||
@patch("json.dump")
|
||||
async def test_remove_config_key(self, mock_json_dump):
|
||||
"""Test removing a key from an object."""
|
||||
self.mock_file_read()
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments(
|
||||
{
|
||||
"operation": "remove",
|
||||
"file_path": self.test_file_path,
|
||||
"json_path": "$.config.enabled",
|
||||
}
|
||||
)
|
||||
)
|
||||
self.assertEqual(result.error_code, 0)
|
||||
|
||||
mock_json_dump.assert_called_once()
|
||||
written_data = mock_json_dump.call_args[0][0]
|
||||
self.assertNotIn("enabled", written_data["config"])
|
||||
|
||||
async def test_view_operation(self):
|
||||
"""Test the view operation to ensure it reads and returns content."""
|
||||
self.mock_file_read()
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments(
|
||||
{
|
||||
"operation": "view",
|
||||
"file_path": self.test_file_path,
|
||||
"json_path": "$.users[0]",
|
||||
}
|
||||
)
|
||||
)
|
||||
self.assertEqual(result.error_code, 0)
|
||||
self.assertIn('"id": 1', result.output)
|
||||
self.assertIn('"name": "Alice"', result.output)
|
||||
|
||||
async def test_error_file_not_found(self):
|
||||
"""Test error handling when the file does not exist."""
|
||||
# Mock Path.exists to return False
|
||||
with (
|
||||
patch("pathlib.Path.exists", return_value=False),
|
||||
patch("pathlib.Path.is_absolute", return_value=True),
|
||||
):
|
||||
result = await self.tool.execute(
|
||||
ToolCallArguments(
|
||||
{
|
||||
"operation": "view",
|
||||
"file_path": "/nonexistent/file.json",
|
||||
}
|
||||
)
|
||||
)
|
||||
self.assertEqual(result.error_code, -1)
|
||||
self.assertIn("File does not exist", result.error)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,76 @@
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from trae_agent.tools.base import ToolCallArguments, ToolExecResult
|
||||
from trae_agent.tools.mcp_tool import MCPTool
|
||||
|
||||
|
||||
class TestMCPTool(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
# simulate a tool schema
|
||||
self.mock_tool = MagicMock()
|
||||
self.mock_tool.name = "test_tool"
|
||||
self.mock_tool.description = "A test tool"
|
||||
self.mock_tool.inputSchema = {
|
||||
"required": ["param1"],
|
||||
"properties": {
|
||||
"param1": {"type": "string", "description": "First parameter"},
|
||||
"param2": {"type": "integer", "description": "Second parameter"},
|
||||
},
|
||||
}
|
||||
|
||||
# simulate client side
|
||||
self.mock_client = MagicMock()
|
||||
self.tool = MCPTool(self.mock_client, self.mock_tool, model_provider="test_provider")
|
||||
|
||||
def test_get_name(self):
|
||||
self.assertEqual(self.tool.get_name(), "test_tool")
|
||||
|
||||
def test_get_description(self):
|
||||
self.assertEqual(self.tool.get_description(), "A test tool")
|
||||
|
||||
def test_get_model_provider(self):
|
||||
self.assertEqual(self.tool.get_model_provider(), "test_provider")
|
||||
|
||||
def test_get_parameters(self):
|
||||
params = self.tool.get_parameters()
|
||||
self.assertEqual(len(params), 2)
|
||||
self.assertTrue(any(p.name == "param1" and p.required for p in params))
|
||||
self.assertTrue(any(p.name == "param2" and not p.required for p in params))
|
||||
|
||||
async def test_execute_success(self):
|
||||
mock_response = MagicMock()
|
||||
mock_response.isError = False
|
||||
mock_response.content = [MagicMock(text="Execution successful")]
|
||||
self.mock_client.call_tool = AsyncMock(return_value=mock_response)
|
||||
|
||||
arguments = ToolCallArguments(arguments={"param1": "value", "param2": 123})
|
||||
result: ToolExecResult = await self.tool.execute(arguments)
|
||||
|
||||
self.assertIsNone(result.error)
|
||||
self.assertEqual(result.output, "Execution successful")
|
||||
|
||||
async def test_execute_failure(self):
|
||||
mock_response = MagicMock()
|
||||
mock_response.isError = True
|
||||
mock_response.content = [MagicMock(text="Something went wrong")]
|
||||
self.mock_client.call_tool = AsyncMock(return_value=mock_response)
|
||||
|
||||
arguments = ToolCallArguments(arguments={"param1": "value"})
|
||||
result: ToolExecResult = await self.tool.execute(arguments)
|
||||
|
||||
self.assertIsNone(result.output)
|
||||
self.assertEqual(result.error, "Something went wrong")
|
||||
|
||||
async def test_execute_exception(self):
|
||||
self.mock_client.call_tool = AsyncMock(side_effect=RuntimeError("Tool crashed"))
|
||||
|
||||
arguments = ToolCallArguments(arguments={"param1": "value"})
|
||||
result: ToolExecResult = await self.tool.execute(arguments)
|
||||
|
||||
self.assertIn("Error running mcp tool", result.error)
|
||||
self.assertEqual(result.error_code, -1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,236 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from trae_agent.utils.config import Config, ModelConfig, ModelProvider
|
||||
from trae_agent.utils.legacy_config import LegacyConfig
|
||||
from trae_agent.utils.llm_clients.anthropic_client import AnthropicClient
|
||||
from trae_agent.utils.llm_clients.openai_client import OpenAIClient
|
||||
|
||||
|
||||
class TestConfigBaseURL(unittest.TestCase):
|
||||
def test_config_with_base_url_in_config(self):
|
||||
test_config = {
|
||||
"default_provider": "openai",
|
||||
"model_providers": {
|
||||
"openai": {
|
||||
"model": "gpt-4o",
|
||||
"api_key": "test-api-key",
|
||||
"base_url": "https://custom-openai.example.com/v1",
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
config = Config.create_from_legacy_config(legacy_config=LegacyConfig(test_config))
|
||||
|
||||
if config.trae_agent:
|
||||
trae_agent_config = config.trae_agent
|
||||
else:
|
||||
self.fail("trae_agent config is None")
|
||||
|
||||
self.assertEqual(
|
||||
trae_agent_config.model.model_provider.base_url,
|
||||
"https://custom-openai.example.com/v1",
|
||||
)
|
||||
|
||||
def test_config_without_base_url(self):
|
||||
test_config = {
|
||||
"default_provider": "openai",
|
||||
"model_providers": {
|
||||
"openai": {
|
||||
"model": "gpt-4o",
|
||||
"api_key": "test-api-key",
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
config = Config.create_from_legacy_config(legacy_config=LegacyConfig(test_config))
|
||||
|
||||
if config.trae_agent:
|
||||
trae_agent_config = config.trae_agent
|
||||
else:
|
||||
self.fail("trae_agent config is None")
|
||||
|
||||
self.assertIsNone(trae_agent_config.model.model_provider.base_url)
|
||||
|
||||
def test_default_anthropic_base_url(self):
|
||||
config = Config.create_from_legacy_config(legacy_config=LegacyConfig({}))
|
||||
|
||||
if config.trae_agent:
|
||||
trae_agent_config = config.trae_agent
|
||||
else:
|
||||
self.fail("trae_agent config is None")
|
||||
|
||||
# If there are no model providers, the default provider is anthropic
|
||||
# and the default base_url is https://api.anthropic.com
|
||||
self.assertEqual(
|
||||
trae_agent_config.model.model_provider.base_url, "https://api.anthropic.com"
|
||||
)
|
||||
|
||||
@patch("trae_agent.utils.llm_clients.openai_client.openai.OpenAI")
|
||||
def test_openai_client_with_custom_base_url(self, mock_openai):
|
||||
model_config = ModelConfig(
|
||||
model="gpt-4o",
|
||||
model_provider=ModelProvider(
|
||||
api_key="test-api-key",
|
||||
provider="openai",
|
||||
base_url="https://custom-openai.example.com/v1",
|
||||
),
|
||||
max_tokens=4096,
|
||||
temperature=0.5,
|
||||
top_p=1,
|
||||
top_k=0,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=10,
|
||||
)
|
||||
|
||||
client = OpenAIClient(model_config)
|
||||
|
||||
mock_openai.assert_called_once_with(
|
||||
api_key="test-api-key", base_url="https://custom-openai.example.com/v1"
|
||||
)
|
||||
self.assertEqual(client.base_url, "https://custom-openai.example.com/v1")
|
||||
|
||||
@patch("trae_agent.utils.llm_clients.anthropic_client.anthropic.Anthropic")
|
||||
def test_anthropic_client_base_url_attribute_set(self, mock_anthropic):
|
||||
model_config = ModelConfig(
|
||||
model="claude-sonnet-4-20250514",
|
||||
model_provider=ModelProvider(
|
||||
api_key="test-api-key",
|
||||
provider="anthropic",
|
||||
base_url="https://custom-anthropic.example.com",
|
||||
),
|
||||
max_tokens=4096,
|
||||
temperature=0.5,
|
||||
top_p=1,
|
||||
top_k=0,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=10,
|
||||
)
|
||||
|
||||
client = AnthropicClient(model_config)
|
||||
|
||||
self.assertEqual(client.base_url, "https://custom-anthropic.example.com")
|
||||
|
||||
@patch("trae_agent.utils.llm_clients.anthropic_client.anthropic.Anthropic")
|
||||
def test_anthropic_client_with_custom_base_url(self, mock_anthropic):
|
||||
model_config = ModelConfig(
|
||||
model="claude-sonnet-4-20250514",
|
||||
model_provider=ModelProvider(
|
||||
api_key="test-api-key",
|
||||
provider="anthropic",
|
||||
base_url="https://custom-anthropic.example.com",
|
||||
),
|
||||
max_tokens=4096,
|
||||
temperature=0.5,
|
||||
top_p=1,
|
||||
top_k=0,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=10,
|
||||
)
|
||||
|
||||
client = AnthropicClient(model_config)
|
||||
|
||||
mock_anthropic.assert_called_once_with(
|
||||
api_key="test-api-key", base_url="https://custom-anthropic.example.com"
|
||||
)
|
||||
self.assertEqual(client.base_url, "https://custom-anthropic.example.com")
|
||||
|
||||
|
||||
class TestLakeviewConfig(unittest.TestCase):
|
||||
def get_base_config(self):
|
||||
return {
|
||||
"default_provider": "anthropic",
|
||||
"enable_lakeview": True,
|
||||
"model_providers": {
|
||||
"anthropic": {
|
||||
"api_key": "anthropic-key",
|
||||
"model": "claude-model",
|
||||
"max_tokens": 4096,
|
||||
"temperature": 0.5,
|
||||
"top_p": 1,
|
||||
"top_k": 0,
|
||||
"max_retries": 10,
|
||||
},
|
||||
"doubao": {
|
||||
"api_key": "doubao-key",
|
||||
"model": "doubao-model",
|
||||
"max_tokens": 8192,
|
||||
"temperature": 0.5,
|
||||
"top_p": 1,
|
||||
"max_retries": 20,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
def get_config_with_mcp_servers(self):
|
||||
return {
|
||||
"default_provider": "anthropic",
|
||||
"enable_lakeview": True,
|
||||
"model_providers": {
|
||||
"anthropic": {
|
||||
"api_key": "anthropic-key",
|
||||
"model": "claude-model",
|
||||
"max_tokens": 4096,
|
||||
"temperature": 0.5,
|
||||
"top_p": 1,
|
||||
"top_k": 0,
|
||||
"max_retries": 10,
|
||||
},
|
||||
"doubao": {
|
||||
"api_key": "doubao-key",
|
||||
"model": "doubao-model",
|
||||
"max_tokens": 8192,
|
||||
"temperature": 0.5,
|
||||
"top_p": 1,
|
||||
"max_retries": 20,
|
||||
},
|
||||
},
|
||||
"mcp_servers": {"test_server": {"command": "echo", "args": [], "env": {}, "cwd": "."}},
|
||||
}
|
||||
|
||||
def test_lakeview_defaults_to_main_provider(self):
|
||||
config_data = self.get_base_config()
|
||||
|
||||
config = Config.create_from_legacy_config(legacy_config=LegacyConfig(config_data))
|
||||
assert config.lakeview is not None
|
||||
self.assertEqual(config.lakeview.model.model_provider.provider, "anthropic")
|
||||
self.assertEqual(config.lakeview.model.model, "claude-model")
|
||||
|
||||
def test_lakeview_null_values_fallback(self):
|
||||
config_data = self.get_base_config()
|
||||
config_data["lakeview_config"] = {"model_provider": None, "model_name": None}
|
||||
|
||||
config = Config.create_from_legacy_config(legacy_config=LegacyConfig(config_data))
|
||||
assert config.lakeview is not None
|
||||
self.assertEqual(config.lakeview.model.model_provider.provider, "anthropic")
|
||||
self.assertEqual(config.lakeview.model.model, "claude-model")
|
||||
|
||||
def test_lakeview_disabled_ignores_config(self):
|
||||
config_data = self.get_base_config()
|
||||
config_data["enable_lakeview"] = False
|
||||
config_data["lakeview_config"] = {"model_provider": "doubao", "model_name": "some-model"}
|
||||
|
||||
config = Config.create_from_legacy_config(legacy_config=LegacyConfig(config_data))
|
||||
self.assertIsNone(config.lakeview)
|
||||
|
||||
def test_mcp_servers_config(self):
|
||||
config_data = self.get_config_with_mcp_servers()
|
||||
config = Config.create_from_legacy_config(legacy_config=LegacyConfig(config_data))
|
||||
self.assertIn("test_server", config.trae_agent.mcp_servers_config)
|
||||
self.assertEqual(config.trae_agent.mcp_servers_config["test_server"].command, "echo")
|
||||
self.assertEqual(config.trae_agent.mcp_servers_config["test_server"].args, [])
|
||||
self.assertEqual(config.trae_agent.mcp_servers_config["test_server"].env, {})
|
||||
self.assertEqual(config.trae_agent.mcp_servers_config["test_server"].cwd, ".")
|
||||
|
||||
def test_mcp_servers_empty_config(self):
|
||||
config_data = self.get_base_config()
|
||||
config = Config.create_from_legacy_config(legacy_config=LegacyConfig(config_data))
|
||||
|
||||
self.assertEqual(config.trae_agent.mcp_servers_config, {})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,315 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""
|
||||
Unit tests for the GoogleClient.
|
||||
|
||||
WARNING: These tests should not be run in a GitHub Actions workflow
|
||||
because they require an API key.
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from trae_agent.tools.base import Tool, ToolCall, ToolResult
|
||||
from trae_agent.utils.config import ModelConfig, ModelProvider
|
||||
from trae_agent.utils.llm_clients.google_client import GoogleClient
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage
|
||||
|
||||
TEST_MODEL = "gemini-2.5-flash"
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
os.getenv("SKIP_GOOGLE_TEST", "").lower() == "true",
|
||||
"Google tests skipped due to SKIP_GOOGLE_TEST environment variable",
|
||||
)
|
||||
class TestGoogleClient(unittest.TestCase):
|
||||
@patch("trae_agent.utils.google_client.genai.Client")
|
||||
def test_google_client_init(self, mock_genai_client):
|
||||
"""Test the initialization of the GoogleClient."""
|
||||
model_config = ModelConfig(
|
||||
model=TEST_MODEL,
|
||||
model_provider=ModelProvider(api_key="test-api-key", provider="google"),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=7.0,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
)
|
||||
google_client = GoogleClient(model_config)
|
||||
mock_genai_client.assert_called_once_with(api_key="test-api-key")
|
||||
self.assertIsNotNone(google_client.client)
|
||||
|
||||
@patch("trae_agent.utils.google_client.genai.Client")
|
||||
@patch.dict(os.environ, {"GOOGLE_API_KEY": "test-env-api-key"})
|
||||
def test_google_client_init_with_env_key(self, mock_genai_client):
|
||||
"""
|
||||
Test that the google client initializes using the GOOGLE_API_KEY environment variable.
|
||||
"""
|
||||
model_config = ModelConfig(
|
||||
model=TEST_MODEL,
|
||||
model_provider=ModelProvider(api_key="", provider="google"),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=7.0,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
)
|
||||
google_client = GoogleClient(model_config)
|
||||
mock_genai_client.assert_called_once_with(api_key="test-env-api-key")
|
||||
self.assertEqual(google_client.api_key, "test-env-api-key")
|
||||
|
||||
@patch.dict(os.environ, {"GOOGLE_API_KEY": ""})
|
||||
def test_google_client_init_no_key_raises_error(self):
|
||||
"""
|
||||
Test that a ValueError is raised if no API key is provided.
|
||||
"""
|
||||
model_config = ModelConfig(
|
||||
model=TEST_MODEL,
|
||||
model_provider=ModelProvider(api_key="", provider="google"),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=7.0,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
)
|
||||
with self.assertRaises(ValueError):
|
||||
GoogleClient(model_config)
|
||||
|
||||
@patch("trae_agent.utils.google_client.genai.Client")
|
||||
def test_google_set_chat_history(self, mock_genai_client):
|
||||
"""
|
||||
Test that the chat history is correctly parsed and stored.
|
||||
"""
|
||||
model_config = ModelConfig(
|
||||
model=TEST_MODEL,
|
||||
model_provider=ModelProvider(api_key="test-api-key", provider="google"),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=7.0,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
)
|
||||
google_client = GoogleClient(model_config)
|
||||
|
||||
messages = [
|
||||
LLMMessage("system", "You are a helpful assistant."),
|
||||
LLMMessage("user", "Hello, world!"),
|
||||
]
|
||||
google_client.set_chat_history(messages)
|
||||
|
||||
self.assertEqual(google_client.system_instruction, "You are a helpful assistant.")
|
||||
self.assertEqual(len(google_client.message_history), 1)
|
||||
self.assertEqual(google_client.message_history[0].role, "user")
|
||||
self.assertEqual(google_client.message_history[0].parts[0].text, "Hello, world!")
|
||||
|
||||
@patch("trae_agent.utils.google_client.genai.Client")
|
||||
def test_google_chat(self, mock_genai_client):
|
||||
"""
|
||||
Test the chat method with a simple user message.
|
||||
"""
|
||||
mock_model = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.candidates = [MagicMock()]
|
||||
mock_response.candidates[0].content.parts = [MagicMock(text="Hello!")]
|
||||
mock_response.candidates[0].finish_reason.name = "STOP"
|
||||
mock_response.usage_metadata = MagicMock(prompt_token_count=10, candidates_token_count=20)
|
||||
mock_model.generate_content.return_value = mock_response
|
||||
mock_genai_client.return_value.models = mock_model
|
||||
|
||||
model_config = ModelConfig(
|
||||
model=TEST_MODEL,
|
||||
model_provider=ModelProvider(api_key="test-api-key", provider="google"),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=7.0,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
)
|
||||
google_client = GoogleClient(model_config)
|
||||
message = LLMMessage("user", "this is a test message")
|
||||
response = google_client.chat(messages=[message], model_config=model_config)
|
||||
|
||||
mock_model.generate_content.assert_called_once()
|
||||
self.assertEqual(response.content, "Hello!")
|
||||
self.assertEqual(response.usage.input_tokens, 10)
|
||||
self.assertEqual(response.usage.output_tokens, 20)
|
||||
self.assertEqual(response.finish_reason, "STOP")
|
||||
|
||||
@patch("trae_agent.utils.google_client.genai.Client")
|
||||
def test_google_chat_with_tool_call(self, mock_genai_client):
|
||||
"""
|
||||
Test the chat method's ability to handle tool calls.
|
||||
"""
|
||||
mock_model = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_function_call = MagicMock()
|
||||
mock_function_call.name = "get_weather"
|
||||
mock_function_call.args = {"location": "Boston"}
|
||||
mock_response.candidates = [MagicMock()]
|
||||
mock_response.candidates[0].content.parts = [
|
||||
MagicMock(function_call=mock_function_call, text=None)
|
||||
]
|
||||
mock_response.candidates[0].finish_reason.name = "TOOL_CALL"
|
||||
mock_response.usage_metadata = MagicMock(prompt_token_count=30, candidates_token_count=15)
|
||||
mock_model.generate_content.return_value = mock_response
|
||||
mock_genai_client.return_value.models = mock_model
|
||||
|
||||
mock_tool = MagicMock(spec=Tool)
|
||||
mock_tool.name = "get_weather"
|
||||
mock_tool.description = "Gets the weather for a location."
|
||||
mock_tool.get_input_schema.return_value = {
|
||||
"type": "object",
|
||||
"properties": {"location": {"type": "string"}},
|
||||
}
|
||||
|
||||
model_config = ModelConfig(
|
||||
model=TEST_MODEL,
|
||||
model_provider=ModelProvider(api_key="test-api-key", provider="google"),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=1.0,
|
||||
top_k=1,
|
||||
parallel_tool_calls=True,
|
||||
max_retries=1,
|
||||
)
|
||||
google_client = GoogleClient(model_config)
|
||||
message = LLMMessage("user", "What is the weather in Boston?")
|
||||
response = google_client.chat(
|
||||
messages=[message], model_config=model_config, tools=[mock_tool]
|
||||
)
|
||||
|
||||
self.assertEqual(response.content, "")
|
||||
self.assertIsNotNone(response.tool_calls)
|
||||
self.assertEqual(len(response.tool_calls), 1)
|
||||
tool_call = response.tool_calls[0]
|
||||
self.assertEqual(tool_call.name, "get_weather")
|
||||
self.assertEqual(tool_call.arguments, {"location": "Boston"})
|
||||
self.assertEqual(response.finish_reason, "TOOL_CALL")
|
||||
|
||||
def test_parse_messages(self):
|
||||
"""Test the parse_messages method with various message types."""
|
||||
google_client = GoogleClient(
|
||||
ModelConfig(
|
||||
model=TEST_MODEL,
|
||||
model_provider=ModelProvider(api_key="test-key", provider="google"),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=1.0,
|
||||
top_k=1,
|
||||
parallel_tool_calls=True,
|
||||
max_retries=1,
|
||||
)
|
||||
)
|
||||
messages = [
|
||||
LLMMessage("system", "Be concise."),
|
||||
LLMMessage("user", "Hello"),
|
||||
LLMMessage(
|
||||
"model",
|
||||
"Hi there!",
|
||||
tool_call=ToolCall(name="search", arguments={"query": "news"}, call_id="tool-123"),
|
||||
),
|
||||
LLMMessage(
|
||||
"tool",
|
||||
"Search results",
|
||||
tool_result=ToolResult(
|
||||
call_id="12345", name="search", result="news data", success=True
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
parsed_messages, system_instruction = google_client.parse_messages(messages)
|
||||
|
||||
self.assertEqual(system_instruction, "Be concise.")
|
||||
self.assertEqual(len(parsed_messages), 3)
|
||||
self.assertEqual(parsed_messages[0].role, "user")
|
||||
self.assertEqual(parsed_messages[0].parts[0].text, "Hello")
|
||||
self.assertEqual(parsed_messages[1].role, "model")
|
||||
self.assertEqual(parsed_messages[1].parts[0].function_call.name, "search")
|
||||
self.assertEqual(parsed_messages[2].role, "tool")
|
||||
self.assertEqual(parsed_messages[2].parts[0].function_response.name, "search")
|
||||
|
||||
def test_parse_tool_call_result(self):
|
||||
"""
|
||||
Test the _parse_tool_call_result method.
|
||||
"""
|
||||
google_client = GoogleClient(
|
||||
ModelConfig(
|
||||
model=TEST_MODEL,
|
||||
model_provider=ModelProvider(api_key="test-key", provider="google"),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=1.0,
|
||||
top_k=1,
|
||||
parallel_tool_calls=True,
|
||||
max_retries=1,
|
||||
)
|
||||
)
|
||||
|
||||
# Test with a simple result
|
||||
tool_result_simple = ToolResult(
|
||||
call_id="1", name="test_tool", result={"status": "done"}, success=True
|
||||
)
|
||||
parsed_part_simple = google_client.parse_tool_call_result(tool_result_simple)
|
||||
self.assertEqual(parsed_part_simple.function_response.name, "test_tool")
|
||||
self.assertEqual(
|
||||
parsed_part_simple.function_response.response,
|
||||
{"result": {"status": "done"}},
|
||||
)
|
||||
|
||||
# Test with an error
|
||||
tool_result_error = ToolResult(
|
||||
call_id="2",
|
||||
name="test_tool",
|
||||
result="some data",
|
||||
error="Something went wrong",
|
||||
success=False,
|
||||
)
|
||||
parsed_part_error = google_client.parse_tool_call_result(tool_result_error)
|
||||
self.assertIn("error", parsed_part_error.function_response.response)
|
||||
self.assertEqual(
|
||||
parsed_part_error.function_response.response["error"],
|
||||
"Something went wrong",
|
||||
)
|
||||
|
||||
# Test with non-serializable result
|
||||
non_serializable_obj = object()
|
||||
tool_result_non_serializable = ToolResult(
|
||||
call_id="3", name="test_tool", result=non_serializable_obj, success=True
|
||||
)
|
||||
parsed_part_non_serializable = google_client.parse_tool_call_result(
|
||||
tool_result_non_serializable
|
||||
)
|
||||
self.assertIn("result", parsed_part_non_serializable.function_response.response)
|
||||
self.assertEqual(
|
||||
parsed_part_non_serializable.function_response.response["result"],
|
||||
str(non_serializable_obj),
|
||||
)
|
||||
|
||||
def test_supports_tool_calling(self):
|
||||
"""Test the supports_tool_calling method."""
|
||||
model_config = ModelConfig(
|
||||
model=TEST_MODEL,
|
||||
model_provider=ModelProvider(api_key="test-api-key", provider="google"),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=7.0,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
base_url=None,
|
||||
)
|
||||
google_client = GoogleClient(model_config)
|
||||
self.assertEqual(google_client.supports_tool_calling(model_config), True)
|
||||
model_config.model = "no such model"
|
||||
self.assertEqual(google_client.supports_tool_calling(model_config), False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,98 @@
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from trae_agent.utils.mcp_client import MCPClient, MCPServerConfig, MCPServerStatus
|
||||
|
||||
|
||||
class TestMCPClient(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.client = MCPClient()
|
||||
|
||||
def test_get_default_server_status(self):
|
||||
status = self.client.get_mcp_server_status("unknown_server")
|
||||
self.assertEqual(status, MCPServerStatus.DISCONNECTED)
|
||||
|
||||
def test_update_and_get_server_status(self):
|
||||
self.client.update_mcp_server_status("test_server", MCPServerStatus.CONNECTED)
|
||||
status = self.client.get_mcp_server_status("test_server")
|
||||
self.assertEqual(status, MCPServerStatus.CONNECTED)
|
||||
|
||||
@patch("trae_agent.utils.mcp_client.ClientSession")
|
||||
async def test_connect_to_server(self, mock_client_session):
|
||||
mock_transport = (MagicMock(), MagicMock())
|
||||
|
||||
mock_instance = mock_client_session.return_value
|
||||
|
||||
mock_instance.initialize = AsyncMock()
|
||||
|
||||
await self.client.connect_to_server("test_server", mock_transport)
|
||||
|
||||
self.assertEqual(
|
||||
self.client.get_mcp_server_status("test_server"), MCPServerStatus.CONNECTED
|
||||
)
|
||||
# mock_instance.initialize.assert_awaited()
|
||||
|
||||
@patch("trae_agent.utils.mcp_client.stdio_client")
|
||||
@patch("trae_agent.utils.mcp_client.ClientSession")
|
||||
async def test_connect_and_discover_stdio(self, mock_client_session, mock_stdio_client):
|
||||
# Setup mock MCP config
|
||||
config = MCPServerConfig(command="echo", args=[], env={}, cwd=".")
|
||||
|
||||
# Mock the returned transport
|
||||
mock_stdio = AsyncMock()
|
||||
mock_writer = AsyncMock()
|
||||
mock_stdio_client.return_value.__aenter__.return_value = (mock_stdio, mock_writer)
|
||||
|
||||
# Mock session and list_tools return
|
||||
mock_session = mock_client_session.return_value
|
||||
mock_session.initialize = AsyncMock()
|
||||
|
||||
mock_session.call_tool = AsyncMock()
|
||||
|
||||
mcp_servers_dict = {}
|
||||
await self.client.connect_and_discover(
|
||||
"test_server", config, mcp_servers_dict, model_provider="mock_provider"
|
||||
)
|
||||
all_tools = []
|
||||
for _, tools in mcp_servers_dict.items():
|
||||
all_tools.extend(tools)
|
||||
self.assertTrue(all(tool.__class__.__name__ == "MCPTool" for tool in all_tools))
|
||||
|
||||
async def test_connect_and_discover_invalid_config(self):
|
||||
config = MCPServerConfig()
|
||||
mcp_servers_dict = {}
|
||||
with self.assertRaises(ValueError):
|
||||
await self.client.connect_and_discover(
|
||||
"invalid_server", config, mcp_servers_dict, model_provider=None
|
||||
)
|
||||
self.assertEqual(len(mcp_servers_dict), 0)
|
||||
|
||||
async def test_call_tool(self):
|
||||
mock_session = AsyncMock()
|
||||
mock_session.call_tool = AsyncMock(return_value={"result": "ok"})
|
||||
self.client.session = mock_session
|
||||
|
||||
result = await self.client.call_tool("tool_name", {"arg1": "val"})
|
||||
self.assertEqual(result, {"result": "ok"})
|
||||
|
||||
async def test_list_tools(self):
|
||||
mock_session = AsyncMock()
|
||||
mock_session.list_tools = AsyncMock(return_value=["tool1", "tool2"])
|
||||
self.client.session = mock_session
|
||||
|
||||
result = await self.client.list_tools()
|
||||
self.assertEqual(result, ["tool1", "tool2"])
|
||||
|
||||
async def test_cleanup(self):
|
||||
self.client.update_mcp_server_status("test_server", MCPServerStatus.CONNECTED)
|
||||
self.client.exit_stack.aclose = AsyncMock()
|
||||
|
||||
await self.client.cleanup("test_server")
|
||||
self.assertEqual(
|
||||
self.client.get_mcp_server_status("test_server"), MCPServerStatus.DISCONNECTED
|
||||
)
|
||||
self.client.exit_stack.aclose.assert_awaited()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,126 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""
|
||||
This test file is used to test the Ollama client. This test program is expected to verify basic functionalities and check if the results match the expected output.
|
||||
|
||||
Currently, we only test init, chat, and set chat history.
|
||||
|
||||
WARNING: This Ollama test should not be used in the GitHub Actions workflow, as using Ollama for testing consumes too much time due to installation.
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
from trae_agent.utils.config import ModelConfig, ModelProvider
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage
|
||||
from trae_agent.utils.llm_clients.ollama_client import OllamaClient
|
||||
|
||||
TEST_MODEL = "qwen3:4b"
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
os.getenv("SKIP_OLLAMA_TEST", "").lower() == "true",
|
||||
"Ollama tests skipped due to SKIP_OLLAMA_TEST environment variable",
|
||||
)
|
||||
class TestOllamaClient(unittest.TestCase):
|
||||
def test_OllamaClient_init(self):
|
||||
"""
|
||||
Test ollama client provides a test case for initialize the ollama client
|
||||
It should not be used to check any configiguration based on BaseLLMClient instead we should just check the parameters
|
||||
that will change during the init process.
|
||||
"""
|
||||
model_config = ModelConfig(
|
||||
TEST_MODEL,
|
||||
model_provider=ModelProvider(
|
||||
provider="ollama",
|
||||
api_key="ollama",
|
||||
base_url="http://localhost:11434/v1",
|
||||
api_version=None,
|
||||
),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=7.0,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
)
|
||||
ollama_client = OllamaClient(model_config)
|
||||
self.assertEqual(ollama_client.api_key, "ollama")
|
||||
self.assertEqual(ollama_client.base_url, "http://localhost:11434/v1")
|
||||
|
||||
def test_ollama_set_chat_history(self):
|
||||
"""
|
||||
There is nothing we have to assert for this test case just see if it can run
|
||||
"""
|
||||
model_config = ModelConfig(
|
||||
TEST_MODEL,
|
||||
model_provider=ModelProvider(
|
||||
provider="ollama",
|
||||
api_key="ollama",
|
||||
base_url="http://localhost:11434/v1",
|
||||
api_version=None,
|
||||
),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=7.0,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
)
|
||||
ollama_client = OllamaClient(model_config)
|
||||
message = LLMMessage("user", "this is a test message")
|
||||
ollama_client.set_chat_history(messages=[message])
|
||||
self.assertTrue(True) # runnable
|
||||
|
||||
def test_ollama_chat(self):
|
||||
"""
|
||||
There is nothing we have to assert for this test case just see if it can run
|
||||
"""
|
||||
model_config = ModelConfig(
|
||||
TEST_MODEL,
|
||||
model_provider=ModelProvider(
|
||||
provider="ollama",
|
||||
api_key="ollama",
|
||||
base_url="http://localhost:11434/v1",
|
||||
api_version=None,
|
||||
),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=7.0,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
)
|
||||
ollama_client = OllamaClient(model_config)
|
||||
message = LLMMessage("user", "this is a test message")
|
||||
ollama_client.chat(messages=[message], model_config=model_config)
|
||||
self.assertTrue(True) # runnable
|
||||
|
||||
def test_supports_tool_calling(self):
|
||||
"""
|
||||
A test case to check the support tool calling function
|
||||
"""
|
||||
model_config = ModelConfig(
|
||||
TEST_MODEL,
|
||||
model_provider=ModelProvider(
|
||||
provider="ollama",
|
||||
api_key="ollama",
|
||||
base_url="http://localhost:11434/v1",
|
||||
api_version=None,
|
||||
),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=7.0,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
)
|
||||
ollama_client = OllamaClient(model_config)
|
||||
self.assertEqual(ollama_client.supports_tool_calling(model_config), True)
|
||||
model_config.model = "no such model"
|
||||
self.assertEqual(ollama_client.supports_tool_calling(model_config), False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,122 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""
|
||||
This file provides basic testing with openrouter client. This purpose of the test is to check if it run properly
|
||||
|
||||
Currently, we only test init, chat and set chat history
|
||||
WARNING: This Open router test should not be used in the GitHub Actions workflow cause it will require API key to test.
|
||||
|
||||
setting: to avoid
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
from trae_agent.utils.config import ModelConfig, ModelProvider
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage
|
||||
from trae_agent.utils.llm_clients.openrouter_client import OpenRouterClient
|
||||
|
||||
TEST_MODEL = "mistralai/mistral-small-3.2-24b-instruct:free"
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
os.getenv("SKIP_OPENROUTER_TEST", "").lower() == "true",
|
||||
"Open router tests skipped due to SKIP_OPENROUTER_TEST environment variable",
|
||||
)
|
||||
class TestOpenRouterClient(unittest.TestCase):
|
||||
"""
|
||||
Open router client init function
|
||||
"""
|
||||
|
||||
def test_OpenRouterClient_init(self):
|
||||
model_config = ModelConfig(
|
||||
TEST_MODEL,
|
||||
model_provider=ModelProvider(
|
||||
provider="openrouter",
|
||||
api_key=os.getenv("OPENROUTER_API_KEY", ""),
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
api_version=None,
|
||||
),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=0.7,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
)
|
||||
openrouter_client = OpenRouterClient(model_config)
|
||||
self.assertEqual(openrouter_client.base_url, "https://openrouter.ai/api/v1")
|
||||
|
||||
def test_set_chat_history(self):
|
||||
model_config = ModelConfig(
|
||||
TEST_MODEL,
|
||||
model_provider=ModelProvider(
|
||||
provider="openrouter",
|
||||
api_key=os.getenv("OPENROUTER_API_KEY", ""),
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
api_version=None,
|
||||
),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=0.7,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
)
|
||||
openrouter_client = OpenRouterClient(model_config)
|
||||
message = LLMMessage("user", "this is a test message")
|
||||
openrouter_client.set_chat_history(messages=[message])
|
||||
self.assertTrue(True) # runnable
|
||||
|
||||
def test_openrouter_chat(self):
|
||||
"""
|
||||
There is nothing we have to assert for this test case just see if it can run
|
||||
"""
|
||||
model_config = ModelConfig(
|
||||
TEST_MODEL,
|
||||
model_provider=ModelProvider(
|
||||
provider="openrouter",
|
||||
api_key=os.getenv("OPENROUTER_API_KEY", ""),
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
api_version=None,
|
||||
),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=0.7,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
)
|
||||
openrouter_client = OpenRouterClient(model_config)
|
||||
message = LLMMessage("user", "this is a test message")
|
||||
openrouter_client.chat(messages=[message], model_config=model_config)
|
||||
self.assertTrue(True) # runnable
|
||||
|
||||
def test_supports_tool_calling(self):
|
||||
"""
|
||||
A test case to check the support tool calling function
|
||||
"""
|
||||
model_config = ModelConfig(
|
||||
model=TEST_MODEL,
|
||||
model_provider=ModelProvider(
|
||||
provider="openrouter",
|
||||
api_key=os.getenv("OPENROUTER_API_KEY", ""),
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
api_version=None,
|
||||
),
|
||||
max_tokens=1000,
|
||||
temperature=0.8,
|
||||
top_p=0.7,
|
||||
top_k=8,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=1,
|
||||
)
|
||||
openrouter_client = OpenRouterClient(model_config)
|
||||
self.assertEqual(openrouter_client.supports_tool_calling(model_config), True)
|
||||
model_config.model = "no such model"
|
||||
self.assertEqual(openrouter_client.supports_tool_calling(model_config), False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,13 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Trae Agent - LLM-based agent for general purpose software engineering tasks."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
|
||||
from trae_agent.agent.base_agent import BaseAgent
|
||||
from trae_agent.agent.trae_agent import TraeAgent
|
||||
from trae_agent.tools.base import Tool, ToolExecutor
|
||||
from trae_agent.utils.llm_clients.llm_client import LLMClient
|
||||
|
||||
__all__ = ["BaseAgent", "TraeAgent", "LLMClient", "Tool", "ToolExecutor"]
|
||||
@@ -0,0 +1,10 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Agent module for Trae Agent."""
|
||||
|
||||
from trae_agent.agent.agent import Agent
|
||||
from trae_agent.agent.base_agent import BaseAgent
|
||||
from trae_agent.agent.trae_agent import TraeAgent
|
||||
|
||||
__all__ = ["BaseAgent", "TraeAgent", "Agent"]
|
||||
@@ -0,0 +1,100 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
from enum import Enum
|
||||
|
||||
from trae_agent.utils.cli.cli_console import CLIConsole
|
||||
from trae_agent.utils.config import AgentConfig, Config
|
||||
from trae_agent.utils.trajectory_recorder import TrajectoryRecorder
|
||||
|
||||
|
||||
class AgentType(Enum):
|
||||
TraeAgent = "trae_agent"
|
||||
|
||||
|
||||
class Agent:
|
||||
def __init__(
|
||||
self,
|
||||
agent_type: AgentType | str,
|
||||
config: Config,
|
||||
trajectory_file: str | None = None,
|
||||
cli_console: CLIConsole | None = None,
|
||||
docker_config: dict | None = None,
|
||||
docker_keep: bool = True,
|
||||
):
|
||||
if isinstance(agent_type, str):
|
||||
agent_type = AgentType(agent_type)
|
||||
self.agent_type: AgentType = agent_type
|
||||
|
||||
# Set up trajectory recording
|
||||
if trajectory_file is not None:
|
||||
self.trajectory_file: str = trajectory_file
|
||||
self.trajectory_recorder: TrajectoryRecorder = TrajectoryRecorder(trajectory_file)
|
||||
else:
|
||||
# Auto-generate trajectory file path
|
||||
self.trajectory_recorder = TrajectoryRecorder()
|
||||
self.trajectory_file = self.trajectory_recorder.get_trajectory_path()
|
||||
|
||||
match self.agent_type:
|
||||
case AgentType.TraeAgent:
|
||||
if config.trae_agent is None:
|
||||
raise ValueError("trae_agent_config is required for TraeAgent")
|
||||
from .trae_agent import TraeAgent
|
||||
|
||||
self.agent_config: AgentConfig = config.trae_agent
|
||||
|
||||
self.agent: TraeAgent = TraeAgent(
|
||||
self.agent_config, docker_config=docker_config, docker_keep=docker_keep
|
||||
)
|
||||
|
||||
self.agent.set_cli_console(cli_console)
|
||||
|
||||
if cli_console:
|
||||
if config.trae_agent.enable_lakeview:
|
||||
cli_console.set_lakeview(config.lakeview)
|
||||
else:
|
||||
cli_console.set_lakeview(None)
|
||||
|
||||
self.agent.set_trajectory_recorder(self.trajectory_recorder)
|
||||
|
||||
async def run(
|
||||
self,
|
||||
task: str,
|
||||
extra_args: dict[str, str] | None = None,
|
||||
tool_names: list[str] | None = None,
|
||||
):
|
||||
self.agent.new_task(task, extra_args, tool_names)
|
||||
|
||||
if self.agent.allow_mcp_servers:
|
||||
if self.agent.cli_console:
|
||||
self.agent.cli_console.print("Initialising MCP tools...")
|
||||
await self.agent.initialise_mcp()
|
||||
|
||||
if self.agent.cli_console:
|
||||
task_details = {
|
||||
"Task": task,
|
||||
"Model Provider": self.agent_config.model.model_provider.provider,
|
||||
"Model": self.agent_config.model.model,
|
||||
"Max Steps": str(self.agent_config.max_steps),
|
||||
"Trajectory File": self.trajectory_file,
|
||||
"Tools": ", ".join([tool.name for tool in self.agent.tools]),
|
||||
}
|
||||
if extra_args:
|
||||
for key, value in extra_args.items():
|
||||
task_details[key.capitalize()] = value
|
||||
self.agent.cli_console.print_task_details(task_details)
|
||||
|
||||
cli_console_task = (
|
||||
asyncio.create_task(self.agent.cli_console.start()) if self.agent.cli_console else None
|
||||
)
|
||||
|
||||
try:
|
||||
execution = await self.agent.execute_task()
|
||||
finally:
|
||||
# Ensure MCP cleanup happens even if execution fails
|
||||
with contextlib.suppress(Exception):
|
||||
await self.agent.cleanup_mcp_clients()
|
||||
|
||||
if cli_console_task:
|
||||
await cli_console_task
|
||||
|
||||
return execution
|
||||
@@ -0,0 +1,100 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
|
||||
from trae_agent.tools.base import ToolCall, ToolResult
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMResponse, LLMUsage
|
||||
|
||||
__all__ = [
|
||||
"AgentStepState",
|
||||
"AgentState",
|
||||
"AgentStep",
|
||||
"AgentExecution",
|
||||
"AgentError",
|
||||
]
|
||||
|
||||
|
||||
class AgentStepState(Enum):
|
||||
"""Defines possible states during an agent's execution lifecycle."""
|
||||
|
||||
THINKING = "thinking"
|
||||
CALLING_TOOL = "calling_tool"
|
||||
REFLECTING = "reflecting"
|
||||
COMPLETED = "completed"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
class AgentState(Enum):
|
||||
"""Defines possible states during an agent's execution lifecycle."""
|
||||
|
||||
IDLE = "idle"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentStep:
|
||||
"""
|
||||
Represents a single step in an agent's execution process.
|
||||
|
||||
Tracks the state, thought process, tool interactions, LLM response,
|
||||
and any associated metadata or errors.
|
||||
"""
|
||||
|
||||
step_number: int
|
||||
state: AgentStepState
|
||||
thought: str | None = None
|
||||
tool_calls: list[ToolCall] | None = None
|
||||
tool_results: list[ToolResult] | None = None
|
||||
llm_response: LLMResponse | None = None
|
||||
reflection: str | None = None
|
||||
error: str | None = None
|
||||
extra: dict[str, object] | None = None
|
||||
llm_usage: LLMUsage | None = None
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"<AgentStep #{self.step_number} "
|
||||
f"state={self.state.name} "
|
||||
f"thought={repr(self.thought)[:40]}...>"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentExecution:
|
||||
"""
|
||||
Encapsulates the entire execution of an agent task.
|
||||
|
||||
Contains the original task, all intermediate steps,
|
||||
final result, execution metadata, and success state.
|
||||
"""
|
||||
|
||||
task: str
|
||||
steps: list[AgentStep]
|
||||
final_result: str | None = None
|
||||
success: bool = False
|
||||
total_tokens: LLMUsage | None = None
|
||||
execution_time: float = 0.0
|
||||
agent_state: AgentState = AgentState.IDLE
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<AgentExecution task={self.task!r} steps={len(self.steps)} success={self.success}>"
|
||||
|
||||
|
||||
class AgentError(Exception):
|
||||
"""
|
||||
Base class for agent-related errors.
|
||||
|
||||
Used to signal execution failures, misconfigurations,
|
||||
or unexpected LLM/tool behavior.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str):
|
||||
self.message: str = message
|
||||
super().__init__(self.message)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<AgentError message={self.message!r}>"
|
||||
@@ -0,0 +1,352 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Base Agent class for LLM-based agents."""
|
||||
|
||||
import contextlib
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Union
|
||||
|
||||
from trae_agent.agent.agent_basics import AgentExecution, AgentState, AgentStep, AgentStepState
|
||||
from trae_agent.agent.docker_manager import DockerManager
|
||||
from trae_agent.tools import tools_registry
|
||||
from trae_agent.tools.base import Tool, ToolCall, ToolExecutor, ToolResult
|
||||
from trae_agent.tools.ckg.ckg_database import clear_older_ckg
|
||||
from trae_agent.tools.docker_tool_executor import DockerToolExecutor
|
||||
from trae_agent.utils.cli import CLIConsole
|
||||
from trae_agent.utils.config import AgentConfig, ModelConfig
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage, LLMResponse
|
||||
from trae_agent.utils.llm_clients.llm_client import LLMClient
|
||||
from trae_agent.utils.trajectory_recorder import TrajectoryRecorder
|
||||
|
||||
|
||||
class BaseAgent(ABC):
|
||||
"""Base class for LLM-based agents."""
|
||||
|
||||
_tool_caller: Union[ToolExecutor, DockerToolExecutor]
|
||||
|
||||
def __init__(
|
||||
self, agent_config: AgentConfig, docker_config: dict | None = None, docker_keep: bool = True
|
||||
):
|
||||
"""Initialize the agent.
|
||||
Args:
|
||||
agent_config: Configuration object containing model parameters and other settings.
|
||||
docker_config: Configuration for running in a Docker environment.
|
||||
"""
|
||||
self._llm_client = LLMClient(agent_config.model)
|
||||
self._model_config = agent_config.model
|
||||
self._max_steps = agent_config.max_steps
|
||||
self._initial_messages: list[LLMMessage] = []
|
||||
self._task: str = ""
|
||||
self._tools: list[Tool] = [
|
||||
tools_registry[tool_name](model_provider=self._model_config.model_provider.provider)
|
||||
for tool_name in agent_config.tools
|
||||
]
|
||||
self.docker_keep = docker_keep
|
||||
self.docker_manager: DockerManager | None = None
|
||||
original_tool_executor = ToolExecutor(self._tools)
|
||||
if docker_config:
|
||||
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
# tools_dir = os.path.join(project_root, 'tools')
|
||||
|
||||
tools_dir = os.path.join(project_root, "dist")
|
||||
|
||||
is_interactive_mode = False
|
||||
self.docker_manager = DockerManager(
|
||||
image=docker_config.get("image"),
|
||||
container_id=docker_config.get("container_id"),
|
||||
dockerfile_path=docker_config.get("dockerfile_path"),
|
||||
docker_image_file=docker_config.get("docker_image_file"),
|
||||
workspace_dir=docker_config["workspace_dir"],
|
||||
tools_dir=tools_dir,
|
||||
interactive=is_interactive_mode,
|
||||
)
|
||||
self._tool_caller = DockerToolExecutor(
|
||||
original_executor=original_tool_executor,
|
||||
docker_manager=self.docker_manager,
|
||||
docker_tools=["bash", "str_replace_based_edit_tool", "json_edit_tool"],
|
||||
host_workspace_dir=docker_config.get("workspace_dir"),
|
||||
container_workspace_dir=self.docker_manager.container_workspace,
|
||||
)
|
||||
else:
|
||||
self._tool_caller = original_tool_executor
|
||||
|
||||
self._cli_console: CLIConsole | None = None
|
||||
|
||||
# Trajectory recorder
|
||||
self._trajectory_recorder: TrajectoryRecorder | None = None
|
||||
|
||||
# CKG tool-specific: clear the older CKG databases
|
||||
clear_older_ckg()
|
||||
|
||||
@property
|
||||
def llm_client(self) -> LLMClient:
|
||||
return self._llm_client
|
||||
|
||||
@property
|
||||
def trajectory_recorder(self) -> TrajectoryRecorder | None:
|
||||
"""Get the trajectory recorder for this agent."""
|
||||
return self._trajectory_recorder
|
||||
|
||||
def set_trajectory_recorder(self, recorder: TrajectoryRecorder | None) -> None:
|
||||
"""Set the trajectory recorder for this agent."""
|
||||
self._trajectory_recorder = recorder
|
||||
# Also set it on the LLM client
|
||||
self._llm_client.set_trajectory_recorder(recorder)
|
||||
|
||||
@property
|
||||
def cli_console(self) -> CLIConsole | None:
|
||||
"""Get the CLI console for this agent."""
|
||||
return self._cli_console
|
||||
|
||||
def set_cli_console(self, cli_console: CLIConsole | None) -> None:
|
||||
"""Set the CLI console for this agent."""
|
||||
self._cli_console = cli_console
|
||||
|
||||
@property
|
||||
def tools(self) -> list[Tool]:
|
||||
"""Get the tools available to this agent."""
|
||||
return self._tools
|
||||
|
||||
@property
|
||||
def task(self) -> str:
|
||||
"""Get the current task of the agent."""
|
||||
return self._task
|
||||
|
||||
@task.setter
|
||||
def task(self, value: str):
|
||||
"""Set the current task of the agent."""
|
||||
self._task = value
|
||||
|
||||
@property
|
||||
def initial_messages(self) -> list[LLMMessage]:
|
||||
"""Get the initial messages for the agent."""
|
||||
return self._initial_messages
|
||||
|
||||
@property
|
||||
def model_config(self) -> ModelConfig:
|
||||
"""Get the model config for the agent."""
|
||||
return self._model_config
|
||||
|
||||
@property
|
||||
def max_steps(self) -> int:
|
||||
"""Get the maximum number of steps for the agent."""
|
||||
return self._max_steps
|
||||
|
||||
@abstractmethod
|
||||
def new_task(
|
||||
self,
|
||||
task: str,
|
||||
extra_args: dict[str, str] | None = None,
|
||||
tool_names: list[str] | None = None,
|
||||
):
|
||||
"""Create a new task."""
|
||||
pass
|
||||
|
||||
async def execute_task(self) -> AgentExecution:
|
||||
"""Execute a task using the agent."""
|
||||
import time
|
||||
|
||||
if self.docker_manager:
|
||||
self.docker_manager.start()
|
||||
|
||||
start_time = time.time()
|
||||
execution = AgentExecution(task=self._task, steps=[])
|
||||
step: AgentStep | None = None
|
||||
|
||||
try:
|
||||
messages = self._initial_messages
|
||||
step_number = 1
|
||||
execution.agent_state = AgentState.RUNNING
|
||||
|
||||
while step_number <= self._max_steps:
|
||||
step = AgentStep(step_number=step_number, state=AgentStepState.THINKING)
|
||||
try:
|
||||
messages = await self._run_llm_step(step, messages, execution)
|
||||
await self._finalize_step(
|
||||
step, messages, execution
|
||||
) # record trajectory for this step and update the CLI console
|
||||
if execution.agent_state == AgentState.COMPLETED:
|
||||
break
|
||||
step_number += 1
|
||||
except Exception as error:
|
||||
execution.agent_state = AgentState.ERROR
|
||||
step.state = AgentStepState.ERROR
|
||||
step.error = str(error)
|
||||
await self._finalize_step(step, messages, execution)
|
||||
break
|
||||
if step_number > self._max_steps and not execution.success:
|
||||
execution.final_result = "Task execution exceeded maximum steps without completion."
|
||||
execution.agent_state = AgentState.ERROR
|
||||
|
||||
except Exception as e:
|
||||
execution.final_result = f"Agent execution failed: {str(e)}"
|
||||
|
||||
finally:
|
||||
if self.docker_manager and not self.docker_keep:
|
||||
self.docker_manager.stop()
|
||||
|
||||
# Ensure tool resources are released whether an exception occurs or not.
|
||||
await self._close_tools()
|
||||
|
||||
execution.execution_time = time.time() - start_time
|
||||
|
||||
# Clean up any MCP clients
|
||||
with contextlib.suppress(Exception):
|
||||
await self.cleanup_mcp_clients()
|
||||
|
||||
self._update_cli_console(step, execution)
|
||||
return execution
|
||||
|
||||
async def _close_tools(self):
|
||||
"""Release tool resources, mainly about BashTool object."""
|
||||
if self._tool_caller:
|
||||
# Ensure all tool resources are properly released.
|
||||
res = await self._tool_caller.close_tools()
|
||||
return res
|
||||
|
||||
async def _run_llm_step(
|
||||
self, step: "AgentStep", messages: list["LLMMessage"], execution: "AgentExecution"
|
||||
) -> list["LLMMessage"]:
|
||||
# Display thinking state
|
||||
step.state = AgentStepState.THINKING
|
||||
self._update_cli_console(step, execution)
|
||||
# Get LLM response
|
||||
llm_response = self._llm_client.chat(messages, self._model_config, self._tools)
|
||||
step.llm_response = llm_response
|
||||
|
||||
# Display step with LLM response
|
||||
self._update_cli_console(step, execution)
|
||||
|
||||
# Update token usage
|
||||
self._update_llm_usage(llm_response, execution)
|
||||
|
||||
if self.llm_indicates_task_completed(llm_response):
|
||||
if self._is_task_completed(llm_response):
|
||||
execution.agent_state = AgentState.COMPLETED
|
||||
execution.final_result = llm_response.content
|
||||
execution.success = True
|
||||
return messages
|
||||
else:
|
||||
execution.agent_state = AgentState.RUNNING
|
||||
return [LLMMessage(role="user", content=self.task_incomplete_message())]
|
||||
else:
|
||||
tool_calls = llm_response.tool_calls
|
||||
return await self._tool_call_handler(tool_calls, step)
|
||||
|
||||
async def _finalize_step(
|
||||
self, step: "AgentStep", messages: list["LLMMessage"], execution: "AgentExecution"
|
||||
) -> None:
|
||||
step.state = AgentStepState.COMPLETED
|
||||
self._record_handler(step, messages)
|
||||
self._update_cli_console(step, execution)
|
||||
execution.steps.append(step)
|
||||
|
||||
def reflect_on_result(self, tool_results: list[ToolResult]) -> str | None:
|
||||
"""Reflect on tool execution result. Override for custom reflection logic."""
|
||||
if len(tool_results) == 0:
|
||||
return None
|
||||
|
||||
reflection = "\n".join(
|
||||
f"The tool execution failed with error: {tool_result.error}. Consider trying a different approach or fixing the parameters."
|
||||
for tool_result in tool_results
|
||||
if not tool_result.success
|
||||
)
|
||||
|
||||
return reflection
|
||||
|
||||
def llm_indicates_task_completed(self, llm_response: LLMResponse) -> bool:
|
||||
"""Check if the LLM indicates that the task is completed. Override for custom logic."""
|
||||
completion_indicators = [
|
||||
"task completed",
|
||||
"task finished",
|
||||
"done",
|
||||
"completed successfully",
|
||||
"finished successfully",
|
||||
]
|
||||
|
||||
response_lower = llm_response.content.lower()
|
||||
return any(indicator in response_lower for indicator in completion_indicators)
|
||||
|
||||
def _is_task_completed(self, llm_response: LLMResponse) -> bool: # pyright: ignore[reportUnusedParameter]
|
||||
"""Check if the task is completed based on the response. Override for custom logic."""
|
||||
return True
|
||||
|
||||
def task_incomplete_message(self) -> str:
|
||||
"""Return a message indicating that the task is incomplete. Override for custom logic."""
|
||||
return "The task is incomplete. Please try again."
|
||||
|
||||
@abstractmethod
|
||||
async def cleanup_mcp_clients(self) -> None:
|
||||
"""Clean up MCP clients. Override in subclasses that use MCP."""
|
||||
pass
|
||||
|
||||
def _update_cli_console(
|
||||
self, step: AgentStep | None = None, agent_execution: AgentExecution | None = None
|
||||
) -> None:
|
||||
if self.cli_console:
|
||||
self.cli_console.update_status(step, agent_execution)
|
||||
|
||||
def _update_llm_usage(self, llm_response: LLMResponse, execution: AgentExecution):
|
||||
if not llm_response.usage:
|
||||
return
|
||||
# if execution.total_tokens is None then set it to be llm_response.usage else sum it up
|
||||
# execution.total_tokens is not None
|
||||
if not execution.total_tokens:
|
||||
execution.total_tokens = llm_response.usage
|
||||
else:
|
||||
execution.total_tokens += llm_response.usage
|
||||
|
||||
def _record_handler(self, step: AgentStep, messages: list[LLMMessage]) -> None:
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_agent_step(
|
||||
step_number=step.step_number,
|
||||
state=step.state.value,
|
||||
llm_messages=messages,
|
||||
llm_response=step.llm_response,
|
||||
tool_calls=step.tool_calls,
|
||||
tool_results=step.tool_results,
|
||||
reflection=step.reflection,
|
||||
error=step.error,
|
||||
)
|
||||
|
||||
async def _tool_call_handler(
|
||||
self, tool_calls: list[ToolCall] | None, step: AgentStep
|
||||
) -> list[LLMMessage]:
|
||||
messages: list[LLMMessage] = []
|
||||
if not tool_calls or len(tool_calls) <= 0:
|
||||
messages = [
|
||||
LLMMessage(
|
||||
role="user",
|
||||
content="It seems that you have not completed the task.",
|
||||
)
|
||||
]
|
||||
return messages
|
||||
|
||||
step.state = AgentStepState.CALLING_TOOL
|
||||
step.tool_calls = tool_calls
|
||||
self._update_cli_console(step)
|
||||
|
||||
if self._model_config.parallel_tool_calls:
|
||||
tool_results = await self._tool_caller.parallel_tool_call(tool_calls)
|
||||
else:
|
||||
tool_results = await self._tool_caller.sequential_tool_call(tool_calls)
|
||||
step.tool_results = tool_results
|
||||
self._update_cli_console(step)
|
||||
for tool_result in tool_results:
|
||||
# Add tool result to conversation
|
||||
message = LLMMessage(role="user", tool_result=tool_result)
|
||||
messages.append(message)
|
||||
|
||||
reflection = self.reflect_on_result(tool_results)
|
||||
if reflection:
|
||||
step.state = AgentStepState.REFLECTING
|
||||
step.reflection = reflection
|
||||
|
||||
# Display reflection
|
||||
self._update_cli_console(step)
|
||||
|
||||
messages.append(LLMMessage(role="assistant", content=reflection))
|
||||
|
||||
return messages
|
||||
@@ -0,0 +1,242 @@
|
||||
import os
|
||||
import subprocess
|
||||
import uuid
|
||||
|
||||
import docker
|
||||
import pexpect
|
||||
from docker.errors import DockerException, ImageNotFound, NotFound
|
||||
|
||||
|
||||
class DockerManager:
|
||||
"""
|
||||
Manages Docker container lifecycle and command execution for the agent.
|
||||
Supports both stateless (non-interactive) and stateful (interactive) modes.
|
||||
"""
|
||||
|
||||
CONTAINER_TOOLS_PATH = "/agent_tools"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
image: str | None,
|
||||
container_id: str | None,
|
||||
dockerfile_path: str | None,
|
||||
docker_image_file: str | None,
|
||||
workspace_dir: str | None = None,
|
||||
tools_dir: str | None = None,
|
||||
interactive: bool = False,
|
||||
):
|
||||
if not image and not container_id and not dockerfile_path and not docker_image_file:
|
||||
raise ValueError(
|
||||
"Either a Docker image or a container ID or a dockerfile path or a docker image file (tar) must be provided."
|
||||
)
|
||||
self.client = docker.from_env()
|
||||
self.image = image
|
||||
self.container_id = container_id
|
||||
self.dockerfile_path = dockerfile_path
|
||||
self.docker_image_file = docker_image_file
|
||||
self.workspace_dir = workspace_dir
|
||||
self.tools_dir = tools_dir
|
||||
self.interactive = interactive
|
||||
self.container_workspace = "/workspace"
|
||||
self.container = None
|
||||
self.shell = None
|
||||
self._is_managed = True
|
||||
|
||||
def start(self):
|
||||
"""Starts/attaches to the container, mounts the workspace, copies tools, and starts the shell."""
|
||||
try:
|
||||
if self.dockerfile_path:
|
||||
if not os.path.isabs(self.dockerfile_path):
|
||||
raise ValueError("Dockerfile path must be an absolute path.")
|
||||
build_context = os.path.dirname(self.dockerfile_path)
|
||||
dockerfile_name = os.path.basename(self.dockerfile_path)
|
||||
unique_tag = f"trae-agent-custom:{uuid.uuid4()}"
|
||||
print(
|
||||
f"Building Docker image from '{self.dockerfile_path}' with tag '{unique_tag}'..."
|
||||
)
|
||||
try:
|
||||
new_image, build_logs = self.client.images.build(
|
||||
path=build_context, dockerfile=dockerfile_name, tag=unique_tag, rm=True
|
||||
)
|
||||
self.image = new_image.tags[0]
|
||||
print(f"✅ Successfully built image: {self.image}")
|
||||
except Exception as e:
|
||||
print("[red]❌ Docker image build failed. See logs below:[/red]")
|
||||
for log_line in e.build_log:
|
||||
if "stream" in log_line:
|
||||
print(log_line["stream"].strip())
|
||||
raise
|
||||
|
||||
elif self.docker_image_file:
|
||||
print(f"Loading Docker image from file '{self.docker_image_file}'...")
|
||||
try:
|
||||
with open(self.docker_image_file, "rb") as f:
|
||||
loaded_images = self.client.images.load(f.read())
|
||||
if not loaded_images:
|
||||
raise DockerException("Failed to load any images from the provided file.")
|
||||
self.image = loaded_images[0].tags[0]
|
||||
print(f"✅ Successfully loaded image: {self.image}")
|
||||
except FileNotFoundError:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise DockerException(f"Error loading image from file: {e}") from e
|
||||
|
||||
if self.container_id:
|
||||
print(f"Attaching to existing container: {self.container_id}...")
|
||||
self.container = self.client.containers.get(self.container_id)
|
||||
self._is_managed = False
|
||||
print(f"Successfully attached to container {self.container.short_id}.")
|
||||
elif self.image:
|
||||
print(f"Starting a new container from image: {self.image}...")
|
||||
if self.workspace_dir is not None:
|
||||
os.makedirs(self.workspace_dir, exist_ok=True)
|
||||
volumes = {
|
||||
os.path.abspath(self.workspace_dir): {
|
||||
"bind": self.container_workspace,
|
||||
"mode": "rw",
|
||||
}
|
||||
}
|
||||
self.container = self.client.containers.run(
|
||||
self.image,
|
||||
command="sleep infinity",
|
||||
detach=True,
|
||||
volumes=volumes,
|
||||
working_dir=self.container_workspace,
|
||||
)
|
||||
self.container_id = self.container.id
|
||||
self._is_managed = True
|
||||
print(
|
||||
f"Container {self.container.short_id} created. Workspace '{self.workspace_dir}' is mounted to '{self.container_workspace}'."
|
||||
)
|
||||
else:
|
||||
self.container = self.client.containers.run(
|
||||
self.image,
|
||||
command="sleep infinity",
|
||||
detach=True,
|
||||
working_dir=self.container_workspace,
|
||||
)
|
||||
self.container_id = self.container.id
|
||||
self._is_managed = True
|
||||
print(f"Container {self.container.short_id} created.")
|
||||
self._copy_tools_to_container()
|
||||
# if self.interactive:
|
||||
self._start_persistent_shell()
|
||||
except (ImageNotFound, NotFound, DockerException) as e:
|
||||
print(f"[red]Failed to start DockerManager: {e}[/red]")
|
||||
raise
|
||||
|
||||
def execute(self, command: str, timeout: int = 300) -> tuple[int, str]:
|
||||
"""
|
||||
Executes a command using the configured mode (interactive or stateless).
|
||||
"""
|
||||
if not self.container:
|
||||
raise RuntimeError("Container is not running. Call start() first.")
|
||||
|
||||
# if self.interactive:
|
||||
return self._execute_interactive(command, timeout)
|
||||
# else:
|
||||
# return self._execute_stateless(command)
|
||||
|
||||
def stop(self):
|
||||
"""Stops the pexpect shell and cleans up the container if managed by this instance."""
|
||||
if self.shell and self.shell.isalive():
|
||||
print("Closing persistent shell...")
|
||||
self.shell.close(force=True)
|
||||
self.shell = None
|
||||
|
||||
if self.container and self._is_managed:
|
||||
print(f"Stopping and removing managed container {self.container.short_id}...")
|
||||
try:
|
||||
self.container.stop()
|
||||
self.container.remove()
|
||||
print("Container cleaned up successfully.")
|
||||
except DockerException as e:
|
||||
print(
|
||||
f"[yellow]Warning: Could not clean up container {self.container.short_id}: {e}[/yellow]"
|
||||
)
|
||||
|
||||
self.container = None
|
||||
|
||||
# --- Private Helper Methods ---
|
||||
|
||||
def _copy_tools_to_container(self):
|
||||
"""Copies the local tools directory to a fixed path inside the container."""
|
||||
if not self.tools_dir or not os.path.isdir(self.tools_dir):
|
||||
print(
|
||||
f"[yellow]Packaged tools directory '{self.tools_dir}' not provided or not found, skipping copy.[/yellow]"
|
||||
)
|
||||
return
|
||||
|
||||
print(
|
||||
f"Copying tools from '{self.tools_dir}' to container path '{self.CONTAINER_TOOLS_PATH}'..."
|
||||
)
|
||||
try:
|
||||
cmd = f"docker cp '{os.path.abspath(self.tools_dir)}' '{self.container.id}:{self.CONTAINER_TOOLS_PATH}'"
|
||||
subprocess.run(cmd, shell=True, check=True, capture_output=True)
|
||||
print("Tools copied successfully.")
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"[red]Failed to copy tools to container: {e.stderr.decode()}[/red]")
|
||||
raise DockerException(f"Failed to copy tools: {e.stderr.decode()}") from e
|
||||
|
||||
def _start_persistent_shell(self):
|
||||
"""Spawns a persistent bash shell inside the container using pexpect."""
|
||||
if not self.container:
|
||||
return
|
||||
# print("Starting persistent shell for interactive mode...")
|
||||
try:
|
||||
command = f"docker exec -it {self.container.id} /bin/bash"
|
||||
self.shell = pexpect.spawn(command, encoding="utf-8", timeout=120)
|
||||
self.shell.expect([r"\$", r"#"], timeout=120)
|
||||
print("Persistent shell is ready.")
|
||||
except pexpect.exceptions.TIMEOUT:
|
||||
print(
|
||||
"[red]Timeout waiting for shell prompt. The container might be slow to start or misconfigured.[/red]"
|
||||
)
|
||||
raise
|
||||
|
||||
# def _execute_stateless(self, command: str) -> tuple[int, str]:
|
||||
# """Executes a command in a new, non-persistent session."""
|
||||
# print(f"Executing (stateless): `{command}`")
|
||||
# exit_code, output_bytes = self.container.exec_run(cmd=f"/bin/sh -c '{command}'")
|
||||
# output = output_bytes.decode('utf-8', errors='replace').strip()
|
||||
# return exit_code, output
|
||||
|
||||
def _execute_interactive(self, command: str, timeout: int) -> tuple[int, str]:
|
||||
"""Executes a command within the existing persistent shell."""
|
||||
if not self.shell or not self.shell.isalive():
|
||||
print("[yellow]Shell not found or died. Attempting to restart...[/yellow]")
|
||||
self._start_persistent_shell()
|
||||
|
||||
if self.shell is None:
|
||||
raise RuntimeError("Failed to start or restart the persistent shell.")
|
||||
|
||||
marker = "---CMD_DONE---"
|
||||
full_command = command.strip()
|
||||
marker_command = f"echo {marker}$?"
|
||||
self.shell.sendline(full_command)
|
||||
self.shell.sendline(marker_command)
|
||||
try:
|
||||
self.shell.expect(marker + r"(\d+)", timeout=timeout)
|
||||
except pexpect.exceptions.TIMEOUT:
|
||||
return (
|
||||
-1,
|
||||
f"Error: Command '{command}' timed out after {timeout} seconds. Partial output:\n{self.shell.before}",
|
||||
)
|
||||
exit_code = int(self.shell.match.group(1))
|
||||
|
||||
output_before_marker = self.shell.before
|
||||
|
||||
# 1. Split the raw output into lines
|
||||
all_lines = output_before_marker.splitlines()
|
||||
# 2. Filter out the lines that are just echoes of our commands
|
||||
clean_lines = []
|
||||
for line in all_lines:
|
||||
stripped_line = line.strip()
|
||||
# Ignore the line if it's an echo of the original command OR our marker command
|
||||
if stripped_line != full_command and marker_command not in stripped_line:
|
||||
clean_lines.append(line)
|
||||
# 3. Join the clean lines back together
|
||||
cleaned_output = "\n".join(clean_lines)
|
||||
# Wait for the next shell prompt to ensure the shell is ready
|
||||
self.shell.expect([r"\$", r"#"])
|
||||
return exit_code, cleaned_output.strip()
|
||||
@@ -0,0 +1,258 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""TraeAgent for software engineering tasks."""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import os
|
||||
import subprocess
|
||||
from typing import override
|
||||
|
||||
from trae_agent.agent.agent_basics import AgentError, AgentExecution
|
||||
from trae_agent.agent.base_agent import BaseAgent
|
||||
from trae_agent.prompt.agent_prompt import TRAE_AGENT_SYSTEM_PROMPT
|
||||
from trae_agent.tools import tools_registry
|
||||
from trae_agent.tools.base import Tool, ToolResult
|
||||
from trae_agent.utils.config import MCPServerConfig, TraeAgentConfig
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage, LLMResponse
|
||||
from trae_agent.utils.mcp_client import MCPClient
|
||||
|
||||
TraeAgentToolNames = [
|
||||
"str_replace_based_edit_tool",
|
||||
"sequentialthinking",
|
||||
"json_edit_tool",
|
||||
"task_done",
|
||||
"bash",
|
||||
]
|
||||
|
||||
|
||||
class TraeAgent(BaseAgent):
|
||||
"""Trae Agent specialized for software engineering tasks."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
trae_agent_config: TraeAgentConfig,
|
||||
docker_config: dict | None = None,
|
||||
docker_keep: bool = True,
|
||||
):
|
||||
"""Initialize TraeAgent.
|
||||
|
||||
Args:
|
||||
config: Configuration object containing model parameters and other settings.
|
||||
Required if llm_client is not provided.
|
||||
llm_client: Optional pre-configured LLMClient instance.
|
||||
If provided, it will be used instead of creating a new one from config.
|
||||
docker_config: Optional configuration for running in a Docker environment.
|
||||
"""
|
||||
self.project_path: str = ""
|
||||
self.base_commit: str | None = None
|
||||
self.must_patch: str = "false"
|
||||
self.patch_path: str | None = None
|
||||
self.mcp_servers_config: dict[str, MCPServerConfig] | None = (
|
||||
trae_agent_config.mcp_servers_config if trae_agent_config.mcp_servers_config else None
|
||||
)
|
||||
self.allow_mcp_servers: list[str] | None = (
|
||||
trae_agent_config.allow_mcp_servers if trae_agent_config.allow_mcp_servers else []
|
||||
)
|
||||
self.mcp_tools: list[Tool] = []
|
||||
self.mcp_clients: list[MCPClient] = [] # Keep track of MCP clients for cleanup
|
||||
self.docker_config = docker_config
|
||||
super().__init__(
|
||||
agent_config=trae_agent_config, docker_config=docker_config, docker_keep=docker_keep
|
||||
)
|
||||
|
||||
async def initialise_mcp(self):
|
||||
"""Async factory to create and initialize TraeAgent."""
|
||||
await self.discover_mcp_tools()
|
||||
|
||||
if self.mcp_tools:
|
||||
self._tools.extend(self.mcp_tools)
|
||||
|
||||
async def discover_mcp_tools(self):
|
||||
if self.mcp_servers_config:
|
||||
for mcp_server_name, mcp_server_config in self.mcp_servers_config.items():
|
||||
if self.allow_mcp_servers is None:
|
||||
return
|
||||
if mcp_server_name not in self.allow_mcp_servers:
|
||||
continue
|
||||
mcp_client = MCPClient()
|
||||
try:
|
||||
await mcp_client.connect_and_discover(
|
||||
mcp_server_name,
|
||||
mcp_server_config,
|
||||
self.mcp_tools,
|
||||
self._llm_client.provider.value,
|
||||
)
|
||||
# Store client for later cleanup
|
||||
self.mcp_clients.append(mcp_client)
|
||||
except Exception:
|
||||
# Clean up failed client
|
||||
with contextlib.suppress(Exception):
|
||||
await mcp_client.cleanup(mcp_server_name)
|
||||
continue
|
||||
except asyncio.CancelledError:
|
||||
# If the task is cancelled, clean up and skip this server
|
||||
with contextlib.suppress(Exception):
|
||||
await mcp_client.cleanup(mcp_server_name)
|
||||
continue
|
||||
else:
|
||||
return
|
||||
|
||||
@override
|
||||
def new_task(
|
||||
self,
|
||||
task: str,
|
||||
extra_args: dict[str, str] | None = None,
|
||||
tool_names: list[str] | None = None,
|
||||
):
|
||||
"""Create a new task."""
|
||||
self._task: str = task
|
||||
|
||||
if tool_names is None and len(self._tools) == 0:
|
||||
tool_names = TraeAgentToolNames
|
||||
|
||||
# Get the model provider from the LLM client
|
||||
provider = self._model_config.model_provider.provider
|
||||
self._tools: list[Tool] = [
|
||||
tools_registry[tool_name](model_provider=provider) for tool_name in tool_names
|
||||
]
|
||||
# self._tool_caller: ToolExecutor = ToolExecutor(self._tools)
|
||||
|
||||
self._initial_messages: list[LLMMessage] = []
|
||||
self._initial_messages.append(LLMMessage(role="system", content=self.get_system_prompt()))
|
||||
|
||||
user_message = ""
|
||||
if not extra_args:
|
||||
raise AgentError("Project path and issue information are required.")
|
||||
if "project_path" not in extra_args:
|
||||
raise AgentError("Project path is required")
|
||||
|
||||
self.project_path = extra_args.get("project_path", "")
|
||||
if self.docker_config:
|
||||
user_message += r"[Project root path]:\workspace\n\n"
|
||||
else:
|
||||
user_message += f"[Project root path]:\n{self.project_path}\n\n"
|
||||
|
||||
if "issue" in extra_args:
|
||||
user_message += f"[Problem statement]: We're currently solving the following issue within our repository. Here's the issue text:\n{extra_args['issue']}\n"
|
||||
optional_attrs_to_set = ["base_commit", "must_patch", "patch_path"]
|
||||
for attr in optional_attrs_to_set:
|
||||
if attr in extra_args:
|
||||
setattr(self, attr, extra_args[attr])
|
||||
|
||||
self._initial_messages.append(LLMMessage(role="user", content=user_message))
|
||||
|
||||
# If trajectory recorder is set, start recording
|
||||
if self._trajectory_recorder:
|
||||
self._trajectory_recorder.start_recording(
|
||||
task=task,
|
||||
provider=self._llm_client.provider.value,
|
||||
model=self._model_config.model,
|
||||
max_steps=self._max_steps,
|
||||
)
|
||||
|
||||
@override
|
||||
async def execute_task(self) -> AgentExecution:
|
||||
"""Execute the task and finalize trajectory recording."""
|
||||
execution = await super().execute_task()
|
||||
|
||||
# Finalize trajectory recording if recorder is available
|
||||
if self._trajectory_recorder:
|
||||
self._trajectory_recorder.finalize_recording(
|
||||
success=execution.success, final_result=execution.final_result
|
||||
)
|
||||
|
||||
if self.patch_path is not None:
|
||||
with open(self.patch_path, "w") as patch_f:
|
||||
_ = patch_f.write(self.get_git_diff())
|
||||
|
||||
return execution
|
||||
|
||||
def get_system_prompt(self) -> str:
|
||||
"""Get the system prompt for TraeAgent."""
|
||||
return TRAE_AGENT_SYSTEM_PROMPT
|
||||
|
||||
@override
|
||||
def reflect_on_result(self, tool_results: list[ToolResult]) -> str | None:
|
||||
return None
|
||||
|
||||
def get_git_diff(self) -> str:
|
||||
"""Get the git diff of the project."""
|
||||
pwd = os.getcwd()
|
||||
if not os.path.isdir(self.project_path):
|
||||
return ""
|
||||
os.chdir(self.project_path)
|
||||
try:
|
||||
if not self.base_commit:
|
||||
stdout = subprocess.check_output(["git", "--no-pager", "diff"]).decode()
|
||||
else:
|
||||
stdout = subprocess.check_output(
|
||||
["git", "--no-pager", "diff", self.base_commit, "HEAD"]
|
||||
).decode()
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
stdout = ""
|
||||
finally:
|
||||
os.chdir(pwd)
|
||||
return stdout
|
||||
|
||||
# Copyright (c) 2024 paul-gauthier
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Original remove_patches_to_tests function was released under Apache-2.0 License, with the full license text
|
||||
# available at https://github.com/Aider-AI/aider-swe-bench/blob/6e98cd6c3b2cbcba12976d6ae1b07f847480cb74/LICENSE.txt
|
||||
# Original function is at https://github.com/Aider-AI/aider-swe-bench/blob/6e98cd6c3b2cbcba12976d6ae1b07f847480cb74/tests.py#L45
|
||||
|
||||
def remove_patches_to_tests(self, model_patch: str) -> str:
|
||||
"""
|
||||
Remove any changes to the tests directory from the provided patch.
|
||||
This is to ensure that the model_patch does not disturb the repo's
|
||||
tests when doing acceptance testing with the `test_patch`.
|
||||
"""
|
||||
lines = model_patch.splitlines(keepends=True)
|
||||
filtered_lines: list[str] = []
|
||||
test_patterns = ["/test/", "/tests/", "/testing/", "test_", "tox.ini"]
|
||||
is_tests = False
|
||||
|
||||
for line in lines:
|
||||
if line.startswith("diff --git a/"):
|
||||
target_path = line.split()[-1]
|
||||
is_tests = target_path.startswith("b/") and any(
|
||||
p in target_path for p in test_patterns
|
||||
)
|
||||
|
||||
if not is_tests:
|
||||
filtered_lines.append(line)
|
||||
|
||||
return "".join(filtered_lines)
|
||||
|
||||
@override
|
||||
def llm_indicates_task_completed(self, llm_response: LLMResponse) -> bool:
|
||||
"""Check if the LLM indicates that the task is completed."""
|
||||
if llm_response.tool_calls is None:
|
||||
return False
|
||||
return any(tool_call.name == "task_done" for tool_call in llm_response.tool_calls)
|
||||
|
||||
@override
|
||||
def _is_task_completed(self, llm_response: LLMResponse) -> bool:
|
||||
"""Enhanced task completion detection."""
|
||||
if self.must_patch == "true":
|
||||
model_patch = self.get_git_diff()
|
||||
patch = self.remove_patches_to_tests(model_patch)
|
||||
if not patch.strip():
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@override
|
||||
def task_incomplete_message(self) -> str:
|
||||
"""Return a message indicating that the task is incomplete."""
|
||||
return "ERROR! Your Patch is empty. Please provide a patch that fixes the problem."
|
||||
|
||||
@override
|
||||
async def cleanup_mcp_clients(self) -> None:
|
||||
"""Clean up all MCP clients to prevent async context leaks."""
|
||||
for client in self.mcp_clients:
|
||||
with contextlib.suppress(Exception):
|
||||
# Use a generic server name for cleanup since we don't track which server each client is for
|
||||
await client.cleanup("cleanup")
|
||||
self.mcp_clients.clear()
|
||||
@@ -0,0 +1,731 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Command Line Interface for Trae Agent."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
|
||||
import click
|
||||
from dotenv import load_dotenv
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.table import Table
|
||||
from rich.text import Text
|
||||
|
||||
from trae_agent.agent import Agent
|
||||
from trae_agent.utils.cli import CLIConsole, ConsoleFactory, ConsoleMode, ConsoleType
|
||||
from trae_agent.utils.config import Config, TraeAgentConfig
|
||||
|
||||
# Load environment variables
|
||||
_ = load_dotenv()
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
def resolve_config_file(config_file: str) -> str:
|
||||
"""
|
||||
Resolve config file with backward compatibility.
|
||||
First tries the specified file, then falls back to JSON if YAML doesn't exist.
|
||||
"""
|
||||
if config_file.endswith(".yaml") or config_file.endswith(".yml"):
|
||||
yaml_path = Path(config_file)
|
||||
json_path = Path(config_file.replace(".yaml", ".json").replace(".yml", ".json"))
|
||||
if yaml_path.exists():
|
||||
return str(yaml_path)
|
||||
elif json_path.exists():
|
||||
console.print(f"[yellow]YAML config not found, using JSON config: {json_path}[/yellow]")
|
||||
return str(json_path)
|
||||
else:
|
||||
console.print(
|
||||
"[red]Error: Config file not found. Please specify a valid config file in the command line option --config-file[/red]"
|
||||
)
|
||||
sys.exit(1)
|
||||
else:
|
||||
return config_file
|
||||
|
||||
|
||||
def check_docker(timeout=3):
|
||||
# 1) Check whether the docker CLI is installed
|
||||
if shutil.which("docker") is None:
|
||||
return {
|
||||
"cli": False,
|
||||
"daemon": False,
|
||||
"version": None,
|
||||
"error": "docker CLI not found",
|
||||
}
|
||||
# 2) Check whether the Docker daemon is reachable (this makes a real request)
|
||||
try:
|
||||
cp = subprocess.run(
|
||||
["docker", "version", "--format", "{{.Server.Version}}"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=timeout,
|
||||
)
|
||||
if cp.returncode == 0 and cp.stdout.strip():
|
||||
return {
|
||||
"cli": True,
|
||||
"daemon": True,
|
||||
"version": cp.stdout.strip(),
|
||||
"error": None,
|
||||
}
|
||||
else:
|
||||
# The daemon may not be running or permissions may be insufficient
|
||||
return {
|
||||
"cli": True,
|
||||
"daemon": False,
|
||||
"version": None,
|
||||
"error": (cp.stderr or cp.stdout).strip(),
|
||||
}
|
||||
except Exception as e:
|
||||
return {"cli": True, "daemon": False, "version": None, "error": str(e)}
|
||||
|
||||
|
||||
def build_with_pyinstaller():
|
||||
os.system("rm -rf trae_agent/dist")
|
||||
print("--- Building edit_tool ---")
|
||||
subprocess.run(
|
||||
[
|
||||
"pyinstaller",
|
||||
"--name",
|
||||
"edit_tool",
|
||||
"trae_agent/tools/edit_tool_cli.py",
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
print("\n--- Building json_edit_tool ---")
|
||||
subprocess.run(
|
||||
[
|
||||
"pyinstaller",
|
||||
"--name",
|
||||
"json_edit_tool",
|
||||
"--hidden-import=jsonpath_ng",
|
||||
"trae_agent/tools/json_edit_tool_cli.py",
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
os.system("mkdir trae_agent/dist")
|
||||
os.system("cp dist/edit_tool/edit_tool trae_agent/dist")
|
||||
os.system("cp -r dist/json_edit_tool/_internal trae_agent/dist")
|
||||
os.system("cp dist/json_edit_tool/json_edit_tool trae_agent/dist")
|
||||
os.system("rm -rf dist")
|
||||
|
||||
|
||||
@click.group()
|
||||
@click.version_option(version="0.1.0")
|
||||
def cli():
|
||||
"""Trae Agent - LLM-based agent for software engineering tasks."""
|
||||
pass
|
||||
|
||||
|
||||
@cli.command()
|
||||
@click.argument("task", required=False)
|
||||
@click.option("--file", "-f", "file_path", help="Path to a file containing the task description.")
|
||||
@click.option("--provider", "-p", help="LLM provider to use")
|
||||
@click.option("--model", "-m", help="Specific model to use")
|
||||
@click.option("--model-base-url", help="Base URL for the model API")
|
||||
@click.option("--api-key", "-k", help="API key (or set via environment variable)")
|
||||
@click.option("--max-steps", help="Maximum number of execution steps", type=int)
|
||||
@click.option("--working-dir", "-w", help="Working directory for the agent")
|
||||
@click.option("--must-patch", "-mp", is_flag=True, help="Whether to patch the code")
|
||||
@click.option(
|
||||
"--config-file",
|
||||
help="Path to configuration file",
|
||||
default="trae_config.yaml",
|
||||
envvar="TRAE_CONFIG_FILE",
|
||||
)
|
||||
@click.option("--trajectory-file", "-t", help="Path to save trajectory file")
|
||||
@click.option("--patch-path", "-pp", help="Path to patch file")
|
||||
# --- Docker Mode Start ---
|
||||
@click.option(
|
||||
"--docker-image",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Specify a Docker image to run the task in a new container",
|
||||
)
|
||||
@click.option(
|
||||
"--docker-container-id",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Attach to an existing Docker container by ID",
|
||||
)
|
||||
@click.option(
|
||||
"--dockerfile-path",
|
||||
type=click.Path(exists=True, dir_okay=False, resolve_path=True),
|
||||
default=None,
|
||||
help="Absolute path to a Dockerfile to build an environment",
|
||||
)
|
||||
@click.option(
|
||||
"--docker-image-file",
|
||||
type=click.Path(exists=True, dir_okay=False, resolve_path=True),
|
||||
default=None,
|
||||
help="Path to a local Docker image file (tar archive) to load.",
|
||||
)
|
||||
@click.option(
|
||||
"--docker-keep",
|
||||
type=bool,
|
||||
default=True,
|
||||
help="Keep or remove the Docker container after finishing the task",
|
||||
)
|
||||
# --- Docker Mode End ---
|
||||
|
||||
|
||||
@click.option(
|
||||
"--console-type",
|
||||
"-ct",
|
||||
default="simple",
|
||||
type=click.Choice(["simple", "rich"], case_sensitive=False),
|
||||
help="Type of console to use (simple or rich)",
|
||||
)
|
||||
@click.option(
|
||||
"--agent-type",
|
||||
"-at",
|
||||
type=click.Choice(["trae_agent"], case_sensitive=False),
|
||||
help="Type of agent to use (trae_agent)",
|
||||
default="trae_agent",
|
||||
)
|
||||
def run(
|
||||
task: str | None,
|
||||
file_path: str | None,
|
||||
patch_path: str,
|
||||
provider: str | None = None,
|
||||
model: str | None = None,
|
||||
model_base_url: str | None = None,
|
||||
api_key: str | None = None,
|
||||
max_steps: int | None = None,
|
||||
working_dir: str | None = None,
|
||||
must_patch: bool = False,
|
||||
config_file: str = "trae_config.yaml",
|
||||
trajectory_file: str | None = None,
|
||||
console_type: str | None = "simple",
|
||||
agent_type: str | None = "trae_agent",
|
||||
# --- Add Docker Mode ---
|
||||
docker_image: str | None = None,
|
||||
docker_container_id: str | None = None,
|
||||
dockerfile_path: str | None = None,
|
||||
docker_image_file: str | None = None,
|
||||
docker_keep: bool = True,
|
||||
):
|
||||
"""
|
||||
Run is the main function of trae. it runs a task using Trae Agent.
|
||||
Args:
|
||||
tasks: the task that you want your agent to solve. This is required to be in the input
|
||||
model: the model expected to be use
|
||||
working_dir: the working directory of the agent. This should be set either in cli or in the config file
|
||||
|
||||
Return:
|
||||
None (it is expected to be ended after calling the run function)
|
||||
"""
|
||||
|
||||
docker_config: dict[str, str | None] | None = None
|
||||
if (
|
||||
sum(
|
||||
[
|
||||
bool(docker_image),
|
||||
bool(docker_container_id),
|
||||
bool(dockerfile_path),
|
||||
bool(docker_image_file),
|
||||
]
|
||||
)
|
||||
> 1
|
||||
):
|
||||
console.print(
|
||||
"[red]Error: --docker-image, --docker-container-id, --dockerfile-path, and --docker-image-file are mutually exclusive.[/red]"
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
if dockerfile_path:
|
||||
docker_config = {"dockerfile_path": dockerfile_path}
|
||||
console.print(
|
||||
f"[blue]Docker mode enabled. Building from Dockerfile: {dockerfile_path}[/blue]"
|
||||
)
|
||||
elif docker_image_file:
|
||||
docker_config = {"docker_image_file": docker_image_file}
|
||||
console.print(
|
||||
f"[blue]Docker mode enabled. Loading from image file: {docker_image_file}[/blue]"
|
||||
)
|
||||
elif docker_container_id:
|
||||
docker_config = {"container_id": docker_container_id}
|
||||
console.print(
|
||||
f"[blue]Docker mode enabled. Attaching to container: {docker_container_id}[/blue]"
|
||||
)
|
||||
elif docker_image:
|
||||
docker_config = {"image": docker_image}
|
||||
console.print(f"[blue]Docker mode enabled. Using image: {docker_image}[/blue]")
|
||||
# --- ADDED END ---
|
||||
|
||||
# Apply backward compatibility for config file
|
||||
config_file = resolve_config_file(config_file)
|
||||
|
||||
if docker_config:
|
||||
check_msg = check_docker()
|
||||
if check_msg["cli"] and check_msg["daemon"] and check_msg["version"]:
|
||||
print("Docker is configured correctly.")
|
||||
else:
|
||||
print(f"Docker is configured incorrectly. {check_msg['error']}")
|
||||
sys.exit(1)
|
||||
if not (os.path.exists("trae_agent/dist") and os.path.exists("trae_agent/dist/_internal")):
|
||||
print("Building tools of Docker mode for the first use, waiting for a few seconds...")
|
||||
build_with_pyinstaller()
|
||||
print("Building finished.")
|
||||
|
||||
if file_path:
|
||||
if task:
|
||||
console.print(
|
||||
"[red]Error: Cannot use both a task string and the --file argument.[/red]"
|
||||
)
|
||||
sys.exit(1)
|
||||
try:
|
||||
task = Path(file_path).read_text()
|
||||
except FileNotFoundError:
|
||||
console.print(f"[red]Error: File not found: {file_path}[/red]")
|
||||
sys.exit(1)
|
||||
elif not task:
|
||||
console.print(
|
||||
"[red]Error: Must provide either a task string or use the --file argument.[/red]"
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
config = Config.create(
|
||||
config_file=config_file,
|
||||
).resolve_config_values(
|
||||
provider=provider,
|
||||
model=model,
|
||||
model_base_url=model_base_url,
|
||||
api_key=api_key,
|
||||
max_steps=max_steps,
|
||||
)
|
||||
|
||||
if not agent_type:
|
||||
console.print("[red]Error: agent_type is required.[/red]")
|
||||
sys.exit(1)
|
||||
|
||||
# Create CLI Console
|
||||
console_mode = ConsoleMode.RUN
|
||||
if console_type:
|
||||
selected_console_type = (
|
||||
ConsoleType.SIMPLE if console_type.lower() == "simple" else ConsoleType.RICH
|
||||
)
|
||||
else:
|
||||
selected_console_type = ConsoleFactory.get_recommended_console_type(console_mode)
|
||||
|
||||
cli_console = ConsoleFactory.create_console(
|
||||
console_type=selected_console_type, mode=console_mode
|
||||
)
|
||||
|
||||
# For rich console in RUN mode, set the initial task
|
||||
if selected_console_type == ConsoleType.RICH and hasattr(cli_console, "set_initial_task"):
|
||||
cli_console.set_initial_task(task)
|
||||
|
||||
# agent = Agent(agent_type, config, trajectory_file, cli_console)
|
||||
|
||||
if docker_config is not None:
|
||||
docker_config["workspace_dir"] = working_dir # now type-safe
|
||||
|
||||
# Change working directory if specified
|
||||
if working_dir:
|
||||
try:
|
||||
Path(working_dir).mkdir(parents=True, exist_ok=True)
|
||||
# os.chdir(working_dir)
|
||||
console.print(f"[blue]Changed working directory to: {working_dir}[/blue]")
|
||||
working_dir = os.path.abspath(working_dir)
|
||||
except Exception as e:
|
||||
error_text = Text(f"Error changing directory: {e}", style="red")
|
||||
console.print(error_text)
|
||||
sys.exit(1)
|
||||
else:
|
||||
working_dir = os.getcwd()
|
||||
console.print(f"[blue]Using current directory as working directory: {working_dir}[/blue]")
|
||||
|
||||
# Ensure working directory is an absolute path
|
||||
if not Path(working_dir).is_absolute():
|
||||
console.print(
|
||||
f"[red]Working directory must be an absolute path: {working_dir}, it should start with `/`[/red]"
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
agent = Agent(
|
||||
agent_type,
|
||||
config,
|
||||
trajectory_file,
|
||||
cli_console,
|
||||
docker_config=docker_config,
|
||||
docker_keep=docker_keep,
|
||||
)
|
||||
|
||||
if not docker_config:
|
||||
try:
|
||||
os.chdir(working_dir)
|
||||
except Exception as e:
|
||||
error_text = Text(f"Error changing directory: {e}", style="red")
|
||||
console.print(error_text)
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
task_args = {
|
||||
"project_path": working_dir,
|
||||
"issue": task,
|
||||
"must_patch": "true" if must_patch else "false",
|
||||
"patch_path": patch_path,
|
||||
}
|
||||
|
||||
# Set up agent context for rich console if applicable
|
||||
if selected_console_type == ConsoleType.RICH and hasattr(cli_console, "set_agent_context"):
|
||||
cli_console.set_agent_context(agent, config.trae_agent, config_file, trajectory_file)
|
||||
|
||||
# Agent will handle starting the appropriate console
|
||||
_ = asyncio.run(agent.run(task, task_args))
|
||||
|
||||
console.print(f"\n[green]Trajectory saved to: {agent.trajectory_file}[/green]")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
console.print("\n[yellow]Task execution interrupted by user[/yellow]")
|
||||
console.print(f"[blue]Partial trajectory saved to: {agent.trajectory_file}[/blue]")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
try:
|
||||
from docker.errors import DockerException
|
||||
|
||||
if isinstance(e, DockerException):
|
||||
error_text = Text(f"Docker Error: {e}", style="red")
|
||||
console.print(f"\n{error_text}")
|
||||
console.print(
|
||||
"[yellow]Please ensure the Docker daemon is running and you have the necessary permissions.[/yellow]"
|
||||
)
|
||||
else:
|
||||
raise e
|
||||
except ImportError:
|
||||
error_text = Text(f"Unexpected error: {e}", style="red")
|
||||
console.print(f"\n{error_text}")
|
||||
console.print(traceback.format_exc())
|
||||
except Exception:
|
||||
error_text = Text(f"Unexpected error: {e}", style="red")
|
||||
console.print(f"\n{error_text}")
|
||||
console.print(traceback.format_exc())
|
||||
console.print(f"[blue]Trajectory saved to: {agent.trajectory_file}[/blue]")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
@cli.command()
|
||||
@click.option("--provider", "-p", help="LLM provider to use")
|
||||
@click.option("--model", "-m", help="Specific model to use")
|
||||
@click.option("--model-base-url", help="Base URL for the model API")
|
||||
@click.option("--api-key", "-k", help="API key (or set via environment variable)")
|
||||
@click.option(
|
||||
"--config-file",
|
||||
help="Path to configuration file",
|
||||
default="trae_config.yaml",
|
||||
envvar="TRAE_CONFIG_FILE",
|
||||
)
|
||||
@click.option("--max-steps", help="Maximum number of execution steps", type=int, default=20)
|
||||
@click.option("--trajectory-file", "-t", help="Path to save trajectory file")
|
||||
@click.option(
|
||||
"--console-type",
|
||||
"-ct",
|
||||
type=click.Choice(["simple", "rich"], case_sensitive=False),
|
||||
help="Type of console to use (simple or rich)",
|
||||
)
|
||||
@click.option(
|
||||
"--agent-type",
|
||||
"-at",
|
||||
type=click.Choice(["trae_agent"], case_sensitive=False),
|
||||
help="Type of agent to use (trae_agent)",
|
||||
default="trae_agent",
|
||||
)
|
||||
def interactive(
|
||||
provider: str | None = None,
|
||||
model: str | None = None,
|
||||
model_base_url: str | None = None,
|
||||
api_key: str | None = None,
|
||||
config_file: str = "trae_config.yaml",
|
||||
max_steps: int | None = None,
|
||||
trajectory_file: str | None = None,
|
||||
console_type: str | None = "simple",
|
||||
agent_type: str | None = "trae_agent",
|
||||
):
|
||||
"""
|
||||
This function starts an interactive session with Trae Agent.
|
||||
Args:
|
||||
console_type: Type of console to use for the interactive session
|
||||
"""
|
||||
# Apply backward compatibility for config file
|
||||
config_file = resolve_config_file(config_file)
|
||||
|
||||
config = Config.create(
|
||||
config_file=config_file,
|
||||
).resolve_config_values(
|
||||
provider=provider,
|
||||
model=model,
|
||||
model_base_url=model_base_url,
|
||||
api_key=api_key,
|
||||
max_steps=max_steps,
|
||||
)
|
||||
|
||||
if config.trae_agent:
|
||||
trae_agent_config = config.trae_agent
|
||||
else:
|
||||
console.print("[red]Error: trae_agent configuration is required in the config file.[/red]")
|
||||
sys.exit(1)
|
||||
|
||||
# Create CLI Console for interactive mode
|
||||
console_mode = ConsoleMode.INTERACTIVE
|
||||
if console_type:
|
||||
selected_console_type = (
|
||||
ConsoleType.SIMPLE if console_type.lower() == "simple" else ConsoleType.RICH
|
||||
)
|
||||
else:
|
||||
selected_console_type = ConsoleFactory.get_recommended_console_type(console_mode)
|
||||
|
||||
cli_console = ConsoleFactory.create_console(
|
||||
console_type=selected_console_type,
|
||||
lakeview_config=config.lakeview,
|
||||
mode=console_mode,
|
||||
)
|
||||
|
||||
if not agent_type:
|
||||
console.print("[red]Error: agent_type is required.[/red]")
|
||||
sys.exit(1)
|
||||
|
||||
# Create agent
|
||||
agent = Agent(agent_type, config, trajectory_file, cli_console)
|
||||
|
||||
# Get the actual trajectory file path (in case it was auto-generated)
|
||||
trajectory_file = agent.trajectory_file
|
||||
|
||||
# For simple console, use traditional interactive loop
|
||||
if selected_console_type == ConsoleType.SIMPLE:
|
||||
asyncio.run(
|
||||
_run_simple_interactive_loop(
|
||||
agent, cli_console, trae_agent_config, config_file, trajectory_file
|
||||
)
|
||||
)
|
||||
else:
|
||||
# For rich console, start the textual app which handles interaction
|
||||
asyncio.run(
|
||||
_run_rich_interactive_loop(
|
||||
agent, cli_console, trae_agent_config, config_file, trajectory_file
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _run_simple_interactive_loop(
|
||||
agent: Agent,
|
||||
cli_console: CLIConsole,
|
||||
trae_agent_config: TraeAgentConfig,
|
||||
config_file: str,
|
||||
trajectory_file: str | None,
|
||||
):
|
||||
"""Run the interactive loop for simple console."""
|
||||
while True:
|
||||
try:
|
||||
task = cli_console.get_task_input()
|
||||
if task is None:
|
||||
console.print("[green]Goodbye![/green]")
|
||||
break
|
||||
|
||||
if task.lower() == "help":
|
||||
console.print(
|
||||
Panel(
|
||||
"""[bold]Available Commands:[/bold]
|
||||
|
||||
• Type any task description to execute it
|
||||
• 'status' - Show agent status
|
||||
• 'clear' - Clear the screen
|
||||
• 'exit' or 'quit' - End the session""",
|
||||
title="Help",
|
||||
border_style="yellow",
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
working_dir = cli_console.get_working_dir_input()
|
||||
|
||||
if task.lower() == "status":
|
||||
console.print(
|
||||
Panel(
|
||||
f"""[bold]Provider:[/bold] {agent.agent_config.model.model_provider.provider}
|
||||
[bold]Model:[/bold] {agent.agent_config.model.model}
|
||||
[bold]Available Tools:[/bold] {len(agent.agent.tools)}
|
||||
[bold]Config File:[/bold] {config_file}
|
||||
[bold]Working Directory:[/bold] {os.getcwd()}""",
|
||||
title="Agent Status",
|
||||
border_style="blue",
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
if task.lower() == "clear":
|
||||
console.clear()
|
||||
continue
|
||||
|
||||
# Set up trajectory recording for this task
|
||||
console.print(f"[blue]Trajectory will be saved to: {trajectory_file}[/blue]")
|
||||
|
||||
task_args = {
|
||||
"project_path": working_dir,
|
||||
"issue": task,
|
||||
"must_patch": "false",
|
||||
}
|
||||
|
||||
# Execute the task
|
||||
console.print(f"\n[blue]Executing task: {task}[/blue]")
|
||||
|
||||
# Start console and execute task
|
||||
console_task = asyncio.create_task(cli_console.start())
|
||||
execution_task = asyncio.create_task(agent.run(task, task_args))
|
||||
|
||||
# Wait for execution to complete
|
||||
_ = await execution_task
|
||||
_ = await console_task
|
||||
|
||||
console.print(f"\n[green]Trajectory saved to: {trajectory_file}[/green]")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
console.print("\n[yellow]Use 'exit' or 'quit' to end the session[/yellow]")
|
||||
except EOFError:
|
||||
console.print("\n[green]Goodbye![/green]")
|
||||
break
|
||||
except Exception as e:
|
||||
error_text = Text(f"Error: {e}", style="red")
|
||||
console.print(error_text)
|
||||
|
||||
|
||||
async def _run_rich_interactive_loop(
|
||||
agent: Agent,
|
||||
cli_console: CLIConsole,
|
||||
trae_agent_config: TraeAgentConfig,
|
||||
config_file: str,
|
||||
trajectory_file: str | None,
|
||||
):
|
||||
"""Run the interactive loop for rich console."""
|
||||
# Set up the agent in the rich console so it can handle task execution
|
||||
if hasattr(cli_console, "set_agent_context"):
|
||||
cli_console.set_agent_context(agent, trae_agent_config, config_file, trajectory_file)
|
||||
|
||||
# Start the console UI - this will handle the entire interaction
|
||||
await cli_console.start()
|
||||
|
||||
|
||||
@cli.command()
|
||||
@click.option(
|
||||
"--config-file",
|
||||
help="Path to configuration file",
|
||||
default="trae_config.yaml",
|
||||
envvar="TRAE_CONFIG_FILE",
|
||||
)
|
||||
@click.option("--provider", "-p", help="LLM provider to use")
|
||||
@click.option("--model", "-m", help="Specific model to use")
|
||||
@click.option("--model-base-url", help="Base URL for the model API")
|
||||
@click.option("--api-key", "-k", help="API key (or set via environment variable)")
|
||||
@click.option("--max-steps", help="Maximum number of execution steps", type=int)
|
||||
def show_config(
|
||||
config_file: str,
|
||||
provider: str | None,
|
||||
model: str | None,
|
||||
model_base_url: str | None,
|
||||
api_key: str | None,
|
||||
max_steps: int | None,
|
||||
):
|
||||
"""Show current configuration settings."""
|
||||
# Apply backward compatibility for config file
|
||||
config_file = resolve_config_file(config_file)
|
||||
|
||||
config_path = Path(config_file)
|
||||
if not config_path.exists():
|
||||
console.print(
|
||||
Panel(
|
||||
f"""[yellow]No configuration file found at: {config_file}[/yellow]
|
||||
|
||||
Using default settings and environment variables.""",
|
||||
title="Configuration Status",
|
||||
border_style="yellow",
|
||||
)
|
||||
)
|
||||
|
||||
config = Config.create(
|
||||
config_file=config_file,
|
||||
).resolve_config_values(
|
||||
provider=provider,
|
||||
model=model,
|
||||
model_base_url=model_base_url,
|
||||
api_key=api_key,
|
||||
max_steps=max_steps,
|
||||
)
|
||||
|
||||
if config.trae_agent:
|
||||
trae_agent_config = config.trae_agent
|
||||
else:
|
||||
console.print("[red]Error: trae_agent configuration is required in the config file.[/red]")
|
||||
sys.exit(1)
|
||||
|
||||
# Display general settings
|
||||
general_table = Table(title="General Settings")
|
||||
general_table.add_column("Setting", style="cyan")
|
||||
general_table.add_column("Value", style="green")
|
||||
|
||||
general_table.add_row(
|
||||
"Default Provider",
|
||||
str(trae_agent_config.model.model_provider.provider or "Not set"),
|
||||
)
|
||||
general_table.add_row("Max Steps", str(trae_agent_config.max_steps or "Not set"))
|
||||
|
||||
console.print(general_table)
|
||||
|
||||
# Display provider settings
|
||||
provider_config = trae_agent_config.model.model_provider
|
||||
provider_table = Table(title=f"{provider_config.provider.title()} Configuration")
|
||||
provider_table.add_column("Setting", style="cyan")
|
||||
provider_table.add_column("Value", style="green")
|
||||
|
||||
provider_table.add_row("Model", trae_agent_config.model.model or "Not set")
|
||||
provider_table.add_row("Base URL", provider_config.base_url or "Not set")
|
||||
provider_table.add_row("API Version", provider_config.api_version or "Not set")
|
||||
provider_table.add_row(
|
||||
"API Key",
|
||||
(
|
||||
f"Set ({provider_config.api_key[:4]}...{provider_config.api_key[-4:]})"
|
||||
if provider_config.api_key
|
||||
else "Not set"
|
||||
),
|
||||
)
|
||||
provider_table.add_row("Max Tokens", str(trae_agent_config.model.max_tokens))
|
||||
provider_table.add_row("Temperature", str(trae_agent_config.model.temperature))
|
||||
provider_table.add_row("Top P", str(trae_agent_config.model.top_p))
|
||||
|
||||
if trae_agent_config.model.model_provider.provider == "anthropic":
|
||||
provider_table.add_row("Top K", str(trae_agent_config.model.top_k))
|
||||
|
||||
console.print(provider_table)
|
||||
|
||||
|
||||
@cli.command()
|
||||
def tools():
|
||||
"""Show available tools and their descriptions."""
|
||||
from .tools import tools_registry
|
||||
|
||||
tools_table = Table(title="Available Tools")
|
||||
tools_table.add_column("Tool Name", style="cyan")
|
||||
tools_table.add_column("Description", style="green")
|
||||
|
||||
for tool_name in tools_registry:
|
||||
try:
|
||||
tool = tools_registry[tool_name]()
|
||||
tools_table.add_row(tool.name, tool.description)
|
||||
except Exception as e:
|
||||
tools_table.add_row(tool_name, f"[red]Error loading: {e}[/red]")
|
||||
|
||||
console.print(tools_table)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point for the CLI."""
|
||||
cli()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,2 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
@@ -0,0 +1,53 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
TRAE_AGENT_SYSTEM_PROMPT = """You are an expert AI software engineering agent.
|
||||
|
||||
File Path Rule: All tools that take a `file_path` as an argument require an **absolute path**. You MUST construct the full, absolute path by combining the `[Project root path]` provided in the user's message with the file's path inside the project.
|
||||
|
||||
For example, if the project root is `/home/user/my_project` and you need to edit `src/main.py`, the correct `file_path` argument is `/home/user/my_project/src/main.py`. Do NOT use relative paths like `src/main.py`.
|
||||
|
||||
Your primary goal is to resolve a given GitHub issue by navigating the provided codebase, identifying the root cause of the bug, implementing a robust fix, and ensuring your changes are safe and well-tested.
|
||||
|
||||
Follow these steps methodically:
|
||||
|
||||
1. Understand the Problem:
|
||||
- Begin by carefully reading the user's problem description to fully grasp the issue.
|
||||
- Identify the core components and expected behavior.
|
||||
|
||||
2. Explore and Locate:
|
||||
- Use the available tools to explore the codebase.
|
||||
- Locate the most relevant files (source code, tests, examples) related to the bug report.
|
||||
|
||||
3. Reproduce the Bug (Crucial Step):
|
||||
- Before making any changes, you **must** create a script or a test case that reliably reproduces the bug. This will be your baseline for verification.
|
||||
- Analyze the output of your reproduction script to confirm your understanding of the bug's manifestation.
|
||||
|
||||
4. Debug and Diagnose:
|
||||
- Inspect the relevant code sections you identified.
|
||||
- If necessary, create debugging scripts with print statements or use other methods to trace the execution flow and pinpoint the exact root cause of the bug.
|
||||
|
||||
5. Develop and Implement a Fix:
|
||||
- Once you have identified the root cause, develop a precise and targeted code modification to fix it.
|
||||
- Use the provided file editing tools to apply your patch. Aim for minimal, clean changes.
|
||||
|
||||
6. Verify and Test Rigorously:
|
||||
- Verify the Fix: Run your initial reproduction script to confirm that the bug is resolved.
|
||||
- Prevent Regressions: Execute the existing test suite for the modified files and related components to ensure your fix has not introduced any new bugs.
|
||||
- Write New Tests: Create new, specific test cases (e.g., using `pytest`) that cover the original bug scenario. This is essential to prevent the bug from recurring in the future. Add these tests to the codebase.
|
||||
- Consider Edge Cases: Think about and test potential edge cases related to your changes.
|
||||
|
||||
7. Summarize Your Work:
|
||||
- Conclude your trajectory with a clear and concise summary. Explain the nature of the bug, the logic of your fix, and the steps you took to verify its correctness and safety.
|
||||
|
||||
**Guiding Principle:** Act like a senior software engineer. Prioritize correctness, safety, and high-quality, test-driven development.
|
||||
|
||||
# GUIDE FOR HOW TO USE "sequential_thinking" TOOL:
|
||||
- Your thinking should be thorough and so it's fine if it's very long. Set total_thoughts to at least 5, but setting it up to 25 is fine as well. You'll need more total thoughts when you are considering multiple possible solutions or root causes for an issue.
|
||||
- Use this tool as much as you find necessary to improve the quality of your answers.
|
||||
- You can run bash commands (like tests, a reproduction script, or 'grep'/'find' to find relevant context) in between thoughts.
|
||||
- The sequential_thinking tool can help you break down complex problems, analyze issues step-by-step, and ensure a thorough approach to problem-solving.
|
||||
- Don't hesitate to use it multiple times throughout your thought process to enhance the depth and accuracy of your solutions.
|
||||
|
||||
If you are sure the issue has been solved, you should call the `task_done` to finish the task.
|
||||
"""
|
||||
@@ -0,0 +1,34 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Tools module for Trae Agent."""
|
||||
|
||||
from trae_agent.tools.base import Tool, ToolCall, ToolExecutor, ToolResult
|
||||
from trae_agent.tools.bash_tool import BashTool
|
||||
from trae_agent.tools.ckg_tool import CKGTool
|
||||
from trae_agent.tools.edit_tool import TextEditorTool
|
||||
from trae_agent.tools.json_edit_tool import JSONEditTool
|
||||
from trae_agent.tools.sequential_thinking_tool import SequentialThinkingTool
|
||||
from trae_agent.tools.task_done_tool import TaskDoneTool
|
||||
|
||||
__all__ = [
|
||||
"Tool",
|
||||
"ToolResult",
|
||||
"ToolCall",
|
||||
"ToolExecutor",
|
||||
"BashTool",
|
||||
"TextEditorTool",
|
||||
"JSONEditTool",
|
||||
"SequentialThinkingTool",
|
||||
"TaskDoneTool",
|
||||
"CKGTool",
|
||||
]
|
||||
|
||||
tools_registry: dict[str, type[Tool]] = {
|
||||
"bash": BashTool,
|
||||
"str_replace_based_edit_tool": TextEditorTool,
|
||||
"json_edit_tool": JSONEditTool,
|
||||
"sequentialthinking": SequentialThinkingTool,
|
||||
"task_done": TaskDoneTool,
|
||||
"ckg": CKGTool,
|
||||
}
|
||||
@@ -0,0 +1,244 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Base classes for tools and tool calling."""
|
||||
|
||||
import asyncio
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from functools import cached_property
|
||||
from typing import TypeAlias, override
|
||||
|
||||
ParamSchemaValue: TypeAlias = str | list[str] | bool | dict[str, object]
|
||||
Property: TypeAlias = dict[str, ParamSchemaValue]
|
||||
|
||||
|
||||
class ToolError(Exception):
|
||||
"""Base class for tool errors."""
|
||||
|
||||
def __init__(self, message: str):
|
||||
super().__init__(message)
|
||||
self.message: str = message
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolExecResult:
|
||||
"""Intermediate result of a tool execution."""
|
||||
|
||||
output: str | None = None
|
||||
error: str | None = None
|
||||
error_code: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolResult:
|
||||
"""Result of a tool execution."""
|
||||
|
||||
call_id: str
|
||||
name: str # Gemini specific field
|
||||
success: bool
|
||||
result: str | None = None
|
||||
error: str | None = None
|
||||
id: str | None = None # OpenAI-specific field
|
||||
|
||||
|
||||
ToolCallArguments = dict[str, str | int | float | dict[str, object] | list[object] | None]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCall:
|
||||
"""Represents a parsed tool call."""
|
||||
|
||||
name: str
|
||||
call_id: str
|
||||
arguments: ToolCallArguments = field(default_factory=dict)
|
||||
id: str | None = None
|
||||
|
||||
@override
|
||||
def __str__(self) -> str:
|
||||
return f"ToolCall(name={self.name}, arguments={self.arguments}, call_id={self.call_id}, id={self.id})"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolParameter:
|
||||
"""Tool parameter definition."""
|
||||
|
||||
name: str
|
||||
type: str | list[str]
|
||||
description: str
|
||||
enum: list[str] | None = None
|
||||
items: dict[str, object] | None = None
|
||||
required: bool = True
|
||||
|
||||
|
||||
class Tool(ABC):
|
||||
"""Base class for all tools."""
|
||||
|
||||
def __init__(self, model_provider: str | None = None):
|
||||
self._model_provider = model_provider
|
||||
|
||||
@cached_property
|
||||
def model_provider(self) -> str | None:
|
||||
return self.get_model_provider()
|
||||
|
||||
@cached_property
|
||||
def name(self) -> str:
|
||||
return self.get_name()
|
||||
|
||||
@cached_property
|
||||
def description(self) -> str:
|
||||
return self.get_description()
|
||||
|
||||
@cached_property
|
||||
def parameters(self) -> list[ToolParameter]:
|
||||
return self.get_parameters()
|
||||
|
||||
def get_model_provider(self) -> str | None:
|
||||
"""Get the model provider."""
|
||||
return self._model_provider
|
||||
|
||||
@abstractmethod
|
||||
def get_name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_parameters(self) -> list[ToolParameter]:
|
||||
"""Get the tool parameters."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def execute(self, arguments: ToolCallArguments) -> ToolExecResult:
|
||||
"""Execute the tool with given parameters."""
|
||||
pass
|
||||
|
||||
def json_definition(self) -> dict[str, object]:
|
||||
return {
|
||||
"name": self.name,
|
||||
"description": self.description,
|
||||
"parameters": self.get_input_schema(),
|
||||
}
|
||||
|
||||
def get_input_schema(self) -> dict[str, object]:
|
||||
"""Get the input schema for the tool."""
|
||||
schema: dict[str, object] = {
|
||||
"type": "object",
|
||||
}
|
||||
|
||||
properties: dict[str, Property] = {}
|
||||
required: list[str] = []
|
||||
|
||||
for param in self.parameters:
|
||||
param_schema: Property = {
|
||||
"type": param.type,
|
||||
"description": param.description,
|
||||
}
|
||||
|
||||
# For OpenAI strict mode, all params must be in 'required'.
|
||||
# Optional params are made "nullable" to be compliant.
|
||||
if self.model_provider == "openai":
|
||||
required.append(param.name)
|
||||
if not param.required:
|
||||
current_type = param_schema["type"]
|
||||
if isinstance(current_type, str):
|
||||
param_schema["type"] = [current_type, "null"]
|
||||
elif isinstance(current_type, list) and "null" not in current_type:
|
||||
param_schema["type"] = list(current_type) + ["null"]
|
||||
elif param.required:
|
||||
required.append(param.name)
|
||||
|
||||
if param.enum:
|
||||
param_schema["enum"] = param.enum
|
||||
|
||||
if param.items:
|
||||
param_schema["items"] = param.items
|
||||
|
||||
# For OpenAI, nested objects also need additionalProperties: false
|
||||
if self.model_provider == "openai" and param.type == "object":
|
||||
param_schema["additionalProperties"] = False
|
||||
|
||||
properties[param.name] = param_schema
|
||||
|
||||
schema["properties"] = properties
|
||||
if len(required) > 0:
|
||||
schema["required"] = required
|
||||
|
||||
# For OpenAI, the top-level schema needs additionalProperties: false
|
||||
if self.model_provider == "openai":
|
||||
schema["additionalProperties"] = False
|
||||
|
||||
return schema
|
||||
|
||||
async def close(self):
|
||||
"""Ensure proper tool resource deallocation before task completion."""
|
||||
return None # Using "pass" will trigger a Ruff check error: B027
|
||||
|
||||
|
||||
class ToolExecutor:
|
||||
"""Tool executor that manages tool execution."""
|
||||
|
||||
def __init__(self, tools: list[Tool]):
|
||||
self._tools = tools
|
||||
self._tool_map: dict[str, Tool] | None = None
|
||||
|
||||
async def close_tools(self):
|
||||
"""Ensure all tool resources are properly released."""
|
||||
tasks = [tool.close() for tool in self._tools if hasattr(tool, "close")]
|
||||
res = await asyncio.gather(*tasks)
|
||||
return res
|
||||
|
||||
def _normalize_name(self, name: str) -> str:
|
||||
"""Normalize tool name by making it lowercase and removing underscores."""
|
||||
return name.lower().replace("_", "")
|
||||
|
||||
@property
|
||||
def tools(self) -> dict[str, Tool]:
|
||||
if self._tool_map is None:
|
||||
self._tool_map = {self._normalize_name(tool.name): tool for tool in self._tools}
|
||||
return self._tool_map
|
||||
|
||||
async def execute_tool_call(self, tool_call: ToolCall) -> ToolResult:
|
||||
"""Execute a tool call."""
|
||||
normalized_name = self._normalize_name(tool_call.name)
|
||||
if normalized_name not in self.tools:
|
||||
return ToolResult(
|
||||
name=tool_call.name,
|
||||
success=False,
|
||||
error=f"Tool '{tool_call.name}' not found. Available tools: {[tool.name for tool in self._tools]}",
|
||||
call_id=tool_call.call_id,
|
||||
id=tool_call.id,
|
||||
)
|
||||
|
||||
tool = self.tools[normalized_name]
|
||||
|
||||
try:
|
||||
tool_exec_result = await tool.execute(tool_call.arguments)
|
||||
return ToolResult(
|
||||
name=tool_call.name,
|
||||
success=tool_exec_result.error_code == 0,
|
||||
result=tool_exec_result.output,
|
||||
error=tool_exec_result.error,
|
||||
call_id=tool_call.call_id,
|
||||
id=tool_call.id,
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolResult(
|
||||
name=tool_call.name,
|
||||
success=False,
|
||||
error=f"Error executing tool '{tool_call.name}': {str(e)}",
|
||||
call_id=tool_call.call_id,
|
||||
id=tool_call.id,
|
||||
)
|
||||
|
||||
async def parallel_tool_call(self, tool_calls: list[ToolCall]) -> list[ToolResult]:
|
||||
"""Execute tool calls in parallel"""
|
||||
return await asyncio.gather(*[self.execute_tool_call(call) for call in tool_calls])
|
||||
|
||||
async def sequential_tool_call(self, tool_calls: list[ToolCall]) -> list[ToolResult]:
|
||||
"""Execute tool calls in sequential"""
|
||||
return [await self.execute_tool_call(call) for call in tool_calls]
|
||||
@@ -0,0 +1,246 @@
|
||||
# Copyright (c) 2023 Anthropic
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: MIT
|
||||
#
|
||||
# This file has been modified by ByteDance Ltd. and/or its affiliates. on 13 June 2025
|
||||
#
|
||||
# Original file was released under MIT License, with the full license text
|
||||
# available at https://github.com/anthropics/anthropic-quickstarts/blob/main/LICENSE
|
||||
#
|
||||
# This modified file is released under the same license.
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from typing import override
|
||||
|
||||
from trae_agent.tools.base import Tool, ToolCallArguments, ToolError, ToolExecResult, ToolParameter
|
||||
|
||||
|
||||
class _BashSession:
|
||||
"""A session of a bash shell."""
|
||||
|
||||
_started: bool
|
||||
_timed_out: bool
|
||||
|
||||
command: str = "/bin/bash"
|
||||
_output_delay: float = 0.2 # seconds
|
||||
_timeout: float = 120.0 # seconds
|
||||
_sentinel: str = ",,,,bash-command-exit-__ERROR_CODE__-banner,,,," # `__ERROR_CODE__` will be replaced by `$?` or `!errorlevel!` later
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._started = False
|
||||
self._timed_out = False
|
||||
self._process: asyncio.subprocess.Process | None = None
|
||||
|
||||
async def start(self) -> None:
|
||||
if self._started:
|
||||
return
|
||||
|
||||
# Windows compatibility: os.setsid not available
|
||||
|
||||
if os.name != "nt": # Unix-like systems
|
||||
self._process = await asyncio.create_subprocess_shell(
|
||||
self.command,
|
||||
shell=True,
|
||||
bufsize=0,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
preexec_fn=os.setsid,
|
||||
)
|
||||
else:
|
||||
self._process = await asyncio.create_subprocess_shell(
|
||||
"cmd.exe /v:on", # enable delayed expansion to allow `echo !errorlevel!`
|
||||
shell=True,
|
||||
bufsize=0,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
|
||||
self._started = True
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Terminate the bash shell."""
|
||||
if not self._started:
|
||||
raise ToolError("Session has not started.")
|
||||
if self._process is None:
|
||||
return
|
||||
if self._process.returncode is not None:
|
||||
return
|
||||
try:
|
||||
self._process.terminate()
|
||||
|
||||
# Wait until the process has truly terminated.
|
||||
stdout, stderr = await asyncio.wait_for(self._process.communicate(), timeout=5.0)
|
||||
except asyncio.TimeoutError:
|
||||
self._process.kill()
|
||||
try:
|
||||
# Set a shorter timeout for the cleanup process
|
||||
stdout, stderr = await asyncio.wait_for(self._process.communicate(), timeout=2.0)
|
||||
except asyncio.TimeoutError:
|
||||
# If it still timeout, return None.
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
async def run(self, command: str) -> ToolExecResult:
|
||||
"""Execute a command in the bash shell."""
|
||||
if not self._started or self._process is None:
|
||||
raise ToolError("Session has not started.")
|
||||
if self._process.returncode is not None:
|
||||
return ToolExecResult(
|
||||
error=f"bash has exited with returncode {self._process.returncode}. tool must be restarted.",
|
||||
error_code=-1,
|
||||
)
|
||||
if self._timed_out:
|
||||
raise ToolError(
|
||||
f"timed out: bash has not returned in {self._timeout} seconds and must be restarted",
|
||||
)
|
||||
|
||||
# we know these are not None because we created the process with PIPEs
|
||||
assert self._process.stdin
|
||||
assert self._process.stdout
|
||||
assert self._process.stderr
|
||||
|
||||
error_code = 0
|
||||
|
||||
sentinel_before, pivot, sentinel_after = self._sentinel.partition("__ERROR_CODE__")
|
||||
assert pivot == "__ERROR_CODE__"
|
||||
|
||||
errcode_retriever = "!errorlevel!" if os.name == "nt" else "$?"
|
||||
command_sep = "&" if os.name == "nt" else ";"
|
||||
|
||||
# send command to the process
|
||||
self._process.stdin.write(
|
||||
b"(\n"
|
||||
+ command.encode()
|
||||
+ f"\n){command_sep} echo {self._sentinel.replace('__ERROR_CODE__', errcode_retriever)}\n".encode()
|
||||
)
|
||||
await self._process.stdin.drain()
|
||||
|
||||
# read output from the process, until the sentinel is found
|
||||
try:
|
||||
async with asyncio.timeout(self._timeout):
|
||||
while True:
|
||||
await asyncio.sleep(self._output_delay)
|
||||
# if we read directly from stdout/stderr, it will wait forever for
|
||||
# EOF. use the StreamReader buffer directly instead.
|
||||
output: str = self._process.stdout._buffer.decode() # type: ignore[attr-defined] # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType]
|
||||
if sentinel_before in output:
|
||||
# strip the sentinel from output
|
||||
output, pivot, exit_banner = output.rpartition(sentinel_before)
|
||||
assert pivot
|
||||
|
||||
# get error code inside banner
|
||||
error_code_str, pivot, _ = exit_banner.partition(sentinel_after)
|
||||
if not pivot or not error_code_str.isdecimal():
|
||||
continue
|
||||
|
||||
error_code = int(error_code_str)
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
self._timed_out = True
|
||||
raise ToolError(
|
||||
f"timed out: bash has not returned in {self._timeout} seconds and must be restarted",
|
||||
) from None
|
||||
|
||||
if output.endswith("\n"): # pyright: ignore[reportUnknownMemberType]
|
||||
output = output[:-1] # pyright: ignore[reportUnknownVariableType]
|
||||
|
||||
error: str = self._process.stderr._buffer.decode() # type: ignore[attr-defined] # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportAttributeAccessIssue]
|
||||
if error.endswith("\n"): # pyright: ignore[reportUnknownMemberType]
|
||||
error = error[:-1] # pyright: ignore[reportUnknownVariableType]
|
||||
|
||||
# clear the buffers so that the next output can be read correctly
|
||||
self._process.stdout._buffer.clear() # type: ignore[attr-defined] # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue]
|
||||
self._process.stderr._buffer.clear() # type: ignore[attr-defined] # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue]
|
||||
|
||||
return ToolExecResult(output=output, error=error, error_code=error_code) # pyright: ignore[reportUnknownArgumentType]
|
||||
|
||||
|
||||
class BashTool(Tool):
|
||||
"""
|
||||
A tool that allows the agent to run bash commands.
|
||||
The tool parameters are defined by Anthropic and are not editable.
|
||||
"""
|
||||
|
||||
def __init__(self, model_provider: str | None = None):
|
||||
super().__init__(model_provider)
|
||||
self._session: _BashSession | None = None
|
||||
|
||||
@override
|
||||
def get_model_provider(self) -> str | None:
|
||||
return self._model_provider
|
||||
|
||||
@override
|
||||
def get_name(self) -> str:
|
||||
return "bash"
|
||||
|
||||
@override
|
||||
def get_description(self) -> str:
|
||||
return """Run commands in a bash shell
|
||||
* When invoking this tool, the contents of the "command" parameter does NOT need to be XML-escaped.
|
||||
* You have access to a mirror of common linux and python packages via apt and pip.
|
||||
* State is persistent across command calls and discussions with the user.
|
||||
* To inspect a particular line range of a file, e.g. lines 10-25, try 'sed -n 10,25p /path/to/the/file'.
|
||||
* Please avoid commands that may produce a very large amount of output.
|
||||
* Please run long lived commands in the background, e.g. 'sleep 10 &' or start a server in the background.
|
||||
"""
|
||||
|
||||
@override
|
||||
def get_parameters(self) -> list[ToolParameter]:
|
||||
# For OpenAI models, all parameters must be required=True
|
||||
# For other providers, optional parameters can have required=False
|
||||
restart_required = self.model_provider == "openai"
|
||||
|
||||
return [
|
||||
ToolParameter(
|
||||
name="command",
|
||||
type="string",
|
||||
description="The bash command to run.",
|
||||
required=True,
|
||||
),
|
||||
ToolParameter(
|
||||
name="restart",
|
||||
type="boolean",
|
||||
description="Set to true to restart the bash session.",
|
||||
required=restart_required,
|
||||
),
|
||||
]
|
||||
|
||||
@override
|
||||
async def execute(self, arguments: ToolCallArguments) -> ToolExecResult:
|
||||
if arguments.get("restart"):
|
||||
if self._session:
|
||||
await self._session.stop()
|
||||
self._session = _BashSession()
|
||||
await self._session.start()
|
||||
|
||||
return ToolExecResult(output="tool has been restarted.")
|
||||
|
||||
if self._session is None:
|
||||
try:
|
||||
self._session = _BashSession()
|
||||
await self._session.start()
|
||||
except Exception as e:
|
||||
return ToolExecResult(error=f"Error starting bash session: {e}", error_code=-1)
|
||||
|
||||
command = str(arguments["command"]) if "command" in arguments else None
|
||||
if command is None:
|
||||
return ToolExecResult(
|
||||
error=f"No command provided for the {self.get_name()} tool",
|
||||
error_code=-1,
|
||||
)
|
||||
try:
|
||||
return await self._session.run(command)
|
||||
except Exception as e:
|
||||
return ToolExecResult(error=f"Error running bash command: {e}", error_code=-1)
|
||||
|
||||
@override
|
||||
async def close(self):
|
||||
"""Properly close self._process."""
|
||||
if self._session:
|
||||
ret = await self._session.stop()
|
||||
self._session = None
|
||||
return ret
|
||||
@@ -0,0 +1,54 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
# Define dataclasses for CKG entries
|
||||
@dataclass
|
||||
class FunctionEntry:
|
||||
"""
|
||||
dataclass for function entry.
|
||||
"""
|
||||
|
||||
name: str
|
||||
file_path: str
|
||||
body: str
|
||||
start_line: int
|
||||
end_line: int
|
||||
parent_function: str | None = None
|
||||
parent_class: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClassEntry:
|
||||
"""
|
||||
dataclass for class entry.
|
||||
"""
|
||||
|
||||
name: str
|
||||
file_path: str
|
||||
body: str
|
||||
start_line: int
|
||||
end_line: int
|
||||
fields: str | None = None
|
||||
methods: str | None = None
|
||||
|
||||
|
||||
# We need a mapping from file extension to tree-sitter language name to parse files and build the graph
|
||||
extension_to_language = {
|
||||
".py": "python",
|
||||
".java": "java",
|
||||
".cpp": "cpp",
|
||||
".hpp": "cpp",
|
||||
".c++": "cpp",
|
||||
".cxx": "cpp",
|
||||
".cc": "cpp",
|
||||
".c": "c",
|
||||
".h": "c",
|
||||
".ts": "typescript",
|
||||
".tsx": "typescript",
|
||||
".js": "javascript",
|
||||
".jsx": "javascript",
|
||||
}
|
||||
@@ -0,0 +1,722 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import sqlite3
|
||||
import subprocess
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from tree_sitter import Node, Parser
|
||||
from tree_sitter_languages import get_parser
|
||||
|
||||
from trae_agent.tools.ckg.base import ClassEntry, FunctionEntry, extension_to_language
|
||||
from trae_agent.utils.constants import LOCAL_STORAGE_PATH
|
||||
|
||||
CKG_DATABASE_PATH = LOCAL_STORAGE_PATH / "ckg"
|
||||
CKG_STORAGE_INFO_FILE = CKG_DATABASE_PATH / "storage_info.json"
|
||||
CKG_DATABASE_EXPIRY_TIME = 60 * 60 * 24 * 7 # 1 week in seconds
|
||||
|
||||
|
||||
"""
|
||||
Known issues:
|
||||
1. When a subdirectory of a codebase that has already been indexed, the CKG is built again for this subdirectory.
|
||||
2. The rebuilding logic can be improved by only rebuilding for files that have been modified.
|
||||
3. For JavaScript and TypeScript, the AST is not complete: anonymous functions, arrow functions, etc., are not parsed.
|
||||
"""
|
||||
|
||||
|
||||
def get_ckg_database_path(codebase_snapshot_hash: str) -> Path:
|
||||
"""Get the path to the CKG database for a codebase path."""
|
||||
return CKG_DATABASE_PATH / f"{codebase_snapshot_hash}.db"
|
||||
|
||||
|
||||
def is_git_repository(folder_path: Path) -> bool:
|
||||
"""Check if the folder is a git repository."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", "rev-parse", "--is-inside-work-tree"],
|
||||
cwd=folder_path,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5,
|
||||
)
|
||||
return result.returncode == 0 and result.stdout.strip() == "true"
|
||||
except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
|
||||
return False
|
||||
|
||||
|
||||
def get_git_status_hash(folder_path: Path) -> str:
|
||||
"""Get hash for git repository (clean or dirty)."""
|
||||
try:
|
||||
# Check if we have any uncommitted changes
|
||||
status_result = subprocess.run(
|
||||
["git", "status", "--porcelain"],
|
||||
cwd=folder_path,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
# Get the current commit hash
|
||||
commit_result = subprocess.run(
|
||||
["git", "rev-parse", "HEAD"], cwd=folder_path, capture_output=True, text=True, timeout=5
|
||||
)
|
||||
|
||||
base_hash = commit_result.stdout.strip()
|
||||
|
||||
# If no uncommitted changes, just use the commit hash
|
||||
if not status_result.stdout.strip():
|
||||
return f"git-clean-{base_hash}"
|
||||
|
||||
# If there are uncommitted changes, include them in the hash
|
||||
uncommitted_hash = hashlib.md5(status_result.stdout.encode()).hexdigest()[:8]
|
||||
return f"git-dirty-{base_hash}-{uncommitted_hash}"
|
||||
|
||||
except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
|
||||
# Fallback to file metadata hash if git commands fail
|
||||
return get_file_metadata_hash(folder_path)
|
||||
|
||||
|
||||
def get_file_metadata_hash(folder_path: Path) -> str:
|
||||
"""Get hash based on file metadata (name, mtime, size) for non-git repositories."""
|
||||
hash_md5 = hashlib.md5()
|
||||
|
||||
for file in folder_path.glob("**/*"):
|
||||
if file.is_file() and not file.name.startswith("."):
|
||||
stat = file.stat()
|
||||
hash_md5.update(file.name.encode())
|
||||
hash_md5.update(str(stat.st_mtime).encode()) # modification time
|
||||
hash_md5.update(str(stat.st_size).encode()) # file size
|
||||
|
||||
return f"metadata-{hash_md5.hexdigest()}"
|
||||
|
||||
|
||||
def get_folder_snapshot_hash(folder_path: Path) -> str:
|
||||
"""Get the hash of the folder snapshot, to make sure that the CKG is up to date."""
|
||||
# Strategy 1: Git repository
|
||||
if is_git_repository(folder_path):
|
||||
return get_git_status_hash(folder_path)
|
||||
|
||||
# Strategy 2: Non-git repository - file metadata
|
||||
return get_file_metadata_hash(folder_path)
|
||||
|
||||
|
||||
def clear_older_ckg():
|
||||
"""Iterate over all the files in the CKG storage directory and delete the ones that are older than 1 week."""
|
||||
for file in CKG_DATABASE_PATH.glob("**/*"):
|
||||
if (
|
||||
file.is_file()
|
||||
and not file.name.startswith(".")
|
||||
and file.name.endswith(".db")
|
||||
and file.stat().st_mtime < datetime.now().timestamp() - CKG_DATABASE_EXPIRY_TIME
|
||||
):
|
||||
try:
|
||||
file.unlink()
|
||||
except Exception as e:
|
||||
print(f"error deleting older CKG database - {file.absolute().as_posix()}: {e}")
|
||||
|
||||
|
||||
SQL_LIST = {
|
||||
"functions": """
|
||||
CREATE TABLE IF NOT EXISTS functions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
name TEXT NOT NULL,
|
||||
file_path TEXT NOT NULL,
|
||||
body TEXT NOT NULL,
|
||||
start_line INTEGER NOT NULL,
|
||||
end_line INTEGER NOT NULL,
|
||||
parent_function TEXT,
|
||||
parent_class TEXT
|
||||
)""",
|
||||
"classes": """
|
||||
CREATE TABLE IF NOT EXISTS classes (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
name TEXT NOT NULL,
|
||||
file_path TEXT NOT NULL,
|
||||
body TEXT NOT NULL,
|
||||
fields TEXT,
|
||||
methods TEXT,
|
||||
start_line INTEGER NOT NULL,
|
||||
end_line INTEGER NOT NULL
|
||||
)""",
|
||||
}
|
||||
|
||||
|
||||
class CKGDatabase:
|
||||
def __init__(self, codebase_path: Path):
|
||||
self._db_connection: sqlite3.Connection
|
||||
self._codebase_path: Path = codebase_path
|
||||
|
||||
if not CKG_DATABASE_PATH.exists():
|
||||
CKG_DATABASE_PATH.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
ckg_storage_info: dict[str, str] = {}
|
||||
|
||||
# to save time and storage, we try to reuse the existing database if the codebase snapshot hash is the same
|
||||
# get the existing codebase snapshot hash from the storage info file
|
||||
if CKG_STORAGE_INFO_FILE.exists():
|
||||
with open(CKG_STORAGE_INFO_FILE, "r") as f:
|
||||
ckg_storage_info = json.load(f)
|
||||
if codebase_path.absolute().as_posix() in ckg_storage_info:
|
||||
existing_codebase_snapshot_hash = ckg_storage_info[
|
||||
codebase_path.absolute().as_posix()
|
||||
]
|
||||
else:
|
||||
existing_codebase_snapshot_hash = ""
|
||||
else:
|
||||
existing_codebase_snapshot_hash = ""
|
||||
|
||||
current_codebase_snapshot_hash = get_folder_snapshot_hash(codebase_path)
|
||||
if existing_codebase_snapshot_hash == current_codebase_snapshot_hash:
|
||||
# we can reuse the existing database
|
||||
database_path = get_ckg_database_path(existing_codebase_snapshot_hash)
|
||||
else:
|
||||
# we need to create a new database and delete the old one
|
||||
database_path = get_ckg_database_path(existing_codebase_snapshot_hash)
|
||||
if database_path.exists():
|
||||
database_path.unlink()
|
||||
database_path = get_ckg_database_path(current_codebase_snapshot_hash)
|
||||
|
||||
ckg_storage_info[codebase_path.absolute().as_posix()] = current_codebase_snapshot_hash
|
||||
with open(CKG_STORAGE_INFO_FILE, "w") as f:
|
||||
json.dump(ckg_storage_info, f)
|
||||
|
||||
if database_path.exists():
|
||||
# reuse existing database
|
||||
self._db_connection = sqlite3.connect(database_path)
|
||||
else:
|
||||
# create new database with tables and build the CKG
|
||||
self._db_connection = sqlite3.connect(database_path)
|
||||
for sql in SQL_LIST.values():
|
||||
self._db_connection.execute(sql)
|
||||
self._db_connection.commit()
|
||||
self._construct_ckg()
|
||||
|
||||
def __del__(self):
|
||||
self._db_connection.close()
|
||||
|
||||
def update(self):
|
||||
"""Update the CKG database."""
|
||||
self._construct_ckg()
|
||||
|
||||
def _recursive_visit_python(
|
||||
self,
|
||||
root_node: Node,
|
||||
file_path: str,
|
||||
parent_class: ClassEntry | None = None,
|
||||
parent_function: FunctionEntry | None = None,
|
||||
):
|
||||
"""Recursively visit the Python AST and insert the entries into the database."""
|
||||
if root_node.type == "function_definition":
|
||||
function_name_node = root_node.child_by_field_name("name")
|
||||
if function_name_node:
|
||||
function_entry = FunctionEntry(
|
||||
name=function_name_node.text.decode(),
|
||||
file_path=file_path,
|
||||
body=root_node.text.decode(),
|
||||
start_line=root_node.start_point[0] + 1,
|
||||
end_line=root_node.end_point[0] + 1,
|
||||
)
|
||||
if parent_function and parent_class:
|
||||
# determine if the function is a method of the class or a function within a function
|
||||
if (
|
||||
parent_function.start_line >= parent_class.start_line
|
||||
and parent_function.end_line <= parent_class.end_line
|
||||
):
|
||||
function_entry.parent_function = parent_function.name
|
||||
else:
|
||||
function_entry.parent_class = parent_class.name
|
||||
elif parent_function:
|
||||
function_entry.parent_function = parent_function.name
|
||||
elif parent_class:
|
||||
function_entry.parent_class = parent_class.name
|
||||
self._insert_entry(function_entry)
|
||||
parent_function = function_entry
|
||||
elif root_node.type == "class_definition":
|
||||
class_name_node = root_node.child_by_field_name("name")
|
||||
if class_name_node:
|
||||
class_body_node = root_node.child_by_field_name("body")
|
||||
class_methods = ""
|
||||
class_entry = ClassEntry(
|
||||
name=class_name_node.text.decode(),
|
||||
file_path=file_path,
|
||||
body=root_node.text.decode(),
|
||||
start_line=root_node.start_point[0] + 1,
|
||||
end_line=root_node.end_point[0] + 1,
|
||||
)
|
||||
if class_body_node:
|
||||
for child in class_body_node.children:
|
||||
function_definition_node = None
|
||||
if child.type == "decorated_definition":
|
||||
function_definition_node = child.child_by_field_name("definition")
|
||||
elif child.type == "function_definition":
|
||||
function_definition_node = child
|
||||
if function_definition_node:
|
||||
method_name_node = function_definition_node.child_by_field_name("name")
|
||||
if method_name_node:
|
||||
parameters_node = function_definition_node.child_by_field_name(
|
||||
"parameters"
|
||||
)
|
||||
return_type_node = child.child_by_field_name("return_type")
|
||||
|
||||
class_method_info = method_name_node.text.decode()
|
||||
if parameters_node:
|
||||
class_method_info += f"{parameters_node.text.decode()}"
|
||||
if return_type_node:
|
||||
class_method_info += f" -> {return_type_node.text.decode()}"
|
||||
class_methods += f"- {class_method_info}\n"
|
||||
class_entry.methods = class_methods.strip() if class_methods != "" else None
|
||||
parent_class = class_entry
|
||||
self._insert_entry(class_entry)
|
||||
|
||||
if len(root_node.children) != 0:
|
||||
for child in root_node.children:
|
||||
self._recursive_visit_python(child, file_path, parent_class, parent_function)
|
||||
|
||||
def _recursive_visit_java(
|
||||
self,
|
||||
root_node: Node,
|
||||
file_path: str,
|
||||
parent_class: ClassEntry | None = None,
|
||||
parent_function: FunctionEntry | None = None,
|
||||
):
|
||||
"""Recursively visit the Java AST and insert the entries into the database."""
|
||||
if root_node.type == "class_declaration":
|
||||
class_name_node = root_node.child_by_field_name("name")
|
||||
if class_name_node:
|
||||
class_entry = ClassEntry(
|
||||
name=class_name_node.text.decode(),
|
||||
file_path=file_path,
|
||||
body=root_node.text.decode(),
|
||||
start_line=root_node.start_point[0] + 1,
|
||||
end_line=root_node.end_point[0] + 1,
|
||||
)
|
||||
class_body_node = root_node.child_by_field_name("body")
|
||||
class_methods = ""
|
||||
class_fields = ""
|
||||
if class_body_node:
|
||||
for child in class_body_node.children:
|
||||
if child.type == "field_declaration":
|
||||
class_fields += f"- {child.text.decode()}\n"
|
||||
if child.type == "method_declaration":
|
||||
method_builder = ""
|
||||
for method_property in child.children:
|
||||
if method_property.type == "block":
|
||||
break
|
||||
method_builder += f"{method_property.text.decode()} "
|
||||
method_builder = method_builder.strip()
|
||||
class_methods += f"- {method_builder}\n"
|
||||
class_entry.methods = class_methods.strip() if class_methods != "" else None
|
||||
class_entry.fields = class_fields.strip() if class_fields != "" else None
|
||||
parent_class = class_entry
|
||||
self._insert_entry(class_entry)
|
||||
elif root_node.type == "method_declaration":
|
||||
method_name_node = root_node.child_by_field_name("name")
|
||||
if method_name_node:
|
||||
method_entry = FunctionEntry(
|
||||
name=method_name_node.text.decode(),
|
||||
file_path=file_path,
|
||||
body=root_node.text.decode(),
|
||||
start_line=root_node.start_point[0] + 1,
|
||||
end_line=root_node.end_point[0] + 1,
|
||||
)
|
||||
if parent_class:
|
||||
method_entry.parent_class = parent_class.name
|
||||
self._insert_entry(method_entry)
|
||||
|
||||
if len(root_node.children) != 0:
|
||||
for child in root_node.children:
|
||||
self._recursive_visit_java(child, file_path, parent_class, parent_function)
|
||||
|
||||
def _recursive_visit_cpp(
|
||||
self,
|
||||
root_node: Node,
|
||||
file_path: str,
|
||||
parent_class: ClassEntry | None = None,
|
||||
parent_function: FunctionEntry | None = None,
|
||||
):
|
||||
"""Recursively visit the C++ AST and insert the entries into the database."""
|
||||
if root_node.type == "class_specifier":
|
||||
class_name_node = root_node.child_by_field_name("name")
|
||||
if class_name_node:
|
||||
class_entry = ClassEntry(
|
||||
name=class_name_node.text.decode(),
|
||||
file_path=file_path,
|
||||
body=root_node.text.decode(),
|
||||
start_line=root_node.start_point[0] + 1,
|
||||
end_line=root_node.end_point[0] + 1,
|
||||
)
|
||||
class_body_node = root_node.child_by_field_name("body")
|
||||
class_methods = ""
|
||||
class_fields = ""
|
||||
if class_body_node:
|
||||
for child in class_body_node.children:
|
||||
if child.type == "function_definition":
|
||||
method_builder = ""
|
||||
for method_property in child.children:
|
||||
if method_property.type == "compound_statement":
|
||||
break
|
||||
method_builder += f"{method_property.text.decode()} "
|
||||
method_builder = method_builder.strip()
|
||||
class_methods += f"- {method_builder}\n"
|
||||
if child.type == "field_declaration":
|
||||
child_is_property = True
|
||||
for child_property in child.children:
|
||||
if child_property.type == "function_declarator":
|
||||
child_is_property = False
|
||||
break
|
||||
if child_is_property:
|
||||
class_fields += f"- {child.text.decode()}\n"
|
||||
else:
|
||||
class_methods += f"- {child.text.decode()}\n"
|
||||
class_entry.methods = class_methods.strip() if class_methods != "" else None
|
||||
class_entry.fields = class_fields.strip() if class_fields != "" else None
|
||||
parent_class = class_entry
|
||||
self._insert_entry(class_entry)
|
||||
elif root_node.type == "function_definition":
|
||||
function_declarator_node = root_node.child_by_field_name("declarator")
|
||||
if function_declarator_node:
|
||||
function_name_node = function_declarator_node.child_by_field_name("declarator")
|
||||
if function_name_node:
|
||||
function_entry = FunctionEntry(
|
||||
name=function_name_node.text.decode(),
|
||||
file_path=file_path,
|
||||
body=root_node.text.decode(),
|
||||
start_line=root_node.start_point[0] + 1,
|
||||
end_line=root_node.end_point[0] + 1,
|
||||
)
|
||||
if parent_class:
|
||||
function_entry.parent_class = parent_class.name
|
||||
self._insert_entry(function_entry)
|
||||
|
||||
if len(root_node.children) != 0:
|
||||
for child in root_node.children:
|
||||
self._recursive_visit_cpp(child, file_path, parent_class, parent_function)
|
||||
|
||||
def _recursive_visit_c(
|
||||
self,
|
||||
root_node: Node,
|
||||
file_path: str,
|
||||
parent_class: ClassEntry | None = None,
|
||||
parent_function: FunctionEntry | None = None,
|
||||
):
|
||||
"""Recursively visit the C AST and insert the entries into the database."""
|
||||
if root_node.type == "function_definition":
|
||||
function_declarator_node = root_node.child_by_field_name("declarator")
|
||||
if function_declarator_node:
|
||||
function_name_node = function_declarator_node.child_by_field_name("declarator")
|
||||
if function_name_node:
|
||||
function_entry = FunctionEntry(
|
||||
name=function_name_node.text.decode(),
|
||||
file_path=file_path,
|
||||
body=root_node.text.decode(),
|
||||
start_line=root_node.start_point[0] + 1,
|
||||
end_line=root_node.end_point[0] + 1,
|
||||
)
|
||||
self._insert_entry(function_entry)
|
||||
|
||||
if len(root_node.children) != 0:
|
||||
for child in root_node.children:
|
||||
self._recursive_visit_c(child, file_path, parent_class, parent_function)
|
||||
|
||||
def _recursive_visit_typescript(
|
||||
self,
|
||||
root_node: Node,
|
||||
file_path: str,
|
||||
parent_class: ClassEntry | None = None,
|
||||
parent_function: FunctionEntry | None = None,
|
||||
):
|
||||
if root_node.type == "class_declaration":
|
||||
class_name_node = root_node.child_by_field_name("name")
|
||||
if class_name_node:
|
||||
class_entry = ClassEntry(
|
||||
name=class_name_node.text.decode(),
|
||||
file_path=file_path,
|
||||
body=root_node.text.decode(),
|
||||
start_line=root_node.start_point[0] + 1,
|
||||
end_line=root_node.end_point[0] + 1,
|
||||
)
|
||||
methods = ""
|
||||
fields = ""
|
||||
class_body_node = root_node.child_by_field_name("body")
|
||||
if class_body_node:
|
||||
for child in class_body_node.children:
|
||||
if child.type == "method_definition":
|
||||
method_builder = ""
|
||||
for method_property in child.children:
|
||||
if method_property.type == "statement_block":
|
||||
break
|
||||
method_builder += f"{method_property.text.decode()} "
|
||||
method_builder = method_builder.strip()
|
||||
methods += f"- {method_builder}\n"
|
||||
elif child.type == "public_field_definition":
|
||||
fields += f"- {child.text.decode()}\n"
|
||||
class_entry.methods = methods.strip() if methods != "" else None
|
||||
class_entry.fields = fields.strip() if fields != "" else None
|
||||
parent_class = class_entry
|
||||
self._insert_entry(class_entry)
|
||||
elif root_node.type == "method_definition":
|
||||
method_name_node = root_node.child_by_field_name("name")
|
||||
if method_name_node:
|
||||
method_entry = FunctionEntry(
|
||||
name=method_name_node.text.decode(),
|
||||
file_path=file_path,
|
||||
body=root_node.text.decode(),
|
||||
start_line=root_node.start_point[0] + 1,
|
||||
end_line=root_node.end_point[0] + 1,
|
||||
)
|
||||
if parent_class:
|
||||
method_entry.parent_class = parent_class.name
|
||||
self._insert_entry(method_entry)
|
||||
|
||||
if len(root_node.children) != 0:
|
||||
for child in root_node.children:
|
||||
self._recursive_visit_typescript(child, file_path, parent_class, parent_function)
|
||||
|
||||
def _recursive_visit_javascript(
|
||||
self,
|
||||
root_node: Node,
|
||||
file_path: str,
|
||||
parent_class: ClassEntry | None = None,
|
||||
parent_function: FunctionEntry | None = None,
|
||||
):
|
||||
"""Recursively visit the JavaScript AST and insert the entries into the database."""
|
||||
if root_node.type == "class_declaration":
|
||||
class_name_node = root_node.child_by_field_name("name")
|
||||
if class_name_node:
|
||||
class_entry = ClassEntry(
|
||||
name=class_name_node.text.decode(),
|
||||
file_path=file_path,
|
||||
body=root_node.text.decode(),
|
||||
start_line=root_node.start_point[0] + 1,
|
||||
end_line=root_node.end_point[0] + 1,
|
||||
)
|
||||
methods = ""
|
||||
fields = ""
|
||||
class_body_node = root_node.child_by_field_name("body")
|
||||
if class_body_node:
|
||||
for child in class_body_node.children:
|
||||
if child.type == "method_definition":
|
||||
method_builder = ""
|
||||
for method_property in child.children:
|
||||
if method_property.type == "statement_block":
|
||||
break
|
||||
method_builder += f"{method_property.text.decode()} "
|
||||
method_builder = method_builder.strip()
|
||||
methods += f"- {method_builder}\n"
|
||||
elif child.type == "public_field_definition":
|
||||
fields += f"- {child.text.decode()}\n"
|
||||
class_entry.methods = methods.strip() if methods != "" else None
|
||||
class_entry.fields = fields.strip() if fields != "" else None
|
||||
parent_class = class_entry
|
||||
self._insert_entry(class_entry)
|
||||
elif root_node.type == "method_definition":
|
||||
method_name_node = root_node.child_by_field_name("name")
|
||||
if method_name_node:
|
||||
method_entry = FunctionEntry(
|
||||
name=method_name_node.text.decode(),
|
||||
file_path=file_path,
|
||||
body=root_node.text.decode(),
|
||||
start_line=root_node.start_point[0] + 1,
|
||||
end_line=root_node.end_point[0] + 1,
|
||||
)
|
||||
if parent_class:
|
||||
method_entry.parent_class = parent_class.name
|
||||
self._insert_entry(method_entry)
|
||||
|
||||
if len(root_node.children) != 0:
|
||||
for child in root_node.children:
|
||||
self._recursive_visit_javascript(child, file_path, parent_class, parent_function)
|
||||
|
||||
def _construct_ckg(self) -> None:
|
||||
"""Initialise the code knowledge graph."""
|
||||
|
||||
# lazy load the parsers for the languages when needed
|
||||
language_to_parser: dict[str, Parser] = {}
|
||||
for file in self._codebase_path.glob("**/*"):
|
||||
# skip hidden files and files in a hidden directory
|
||||
if (
|
||||
file.is_file()
|
||||
and not file.name.startswith(".")
|
||||
and "/." not in file.absolute().as_posix()
|
||||
):
|
||||
extension = file.suffix
|
||||
# ignore files with unknown extensions
|
||||
if extension not in extension_to_language:
|
||||
continue
|
||||
language = extension_to_language[extension]
|
||||
|
||||
language_parser = language_to_parser.get(language)
|
||||
if not language_parser:
|
||||
language_parser = get_parser(language)
|
||||
language_to_parser[language] = language_parser
|
||||
|
||||
tree = language_parser.parse(file.read_bytes())
|
||||
root_node = tree.root_node
|
||||
|
||||
match language:
|
||||
case "python":
|
||||
self._recursive_visit_python(root_node, file.absolute().as_posix())
|
||||
case "java":
|
||||
self._recursive_visit_java(root_node, file.absolute().as_posix())
|
||||
case "cpp":
|
||||
self._recursive_visit_cpp(root_node, file.absolute().as_posix())
|
||||
case "c":
|
||||
self._recursive_visit_c(root_node, file.absolute().as_posix())
|
||||
case "typescript":
|
||||
self._recursive_visit_typescript(root_node, file.absolute().as_posix())
|
||||
case "javascript":
|
||||
self._recursive_visit_javascript(root_node, file.absolute().as_posix())
|
||||
case _:
|
||||
continue
|
||||
|
||||
def _insert_entry(self, entry: FunctionEntry | ClassEntry) -> None:
|
||||
"""
|
||||
Insert entry into db.
|
||||
|
||||
Args:
|
||||
entry: the entry to insert
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# TODO: add try catch block to avoid connection problem.
|
||||
match entry:
|
||||
case FunctionEntry():
|
||||
self._insert_function(entry)
|
||||
|
||||
case ClassEntry():
|
||||
self._insert_class(entry)
|
||||
|
||||
self._db_connection.commit()
|
||||
|
||||
def _insert_function(self, entry: FunctionEntry) -> None:
|
||||
"""
|
||||
Insert function entry including functions and class methodsinto db.
|
||||
|
||||
Args:
|
||||
entry: the entry to insert
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
self._db_connection.execute(
|
||||
"""
|
||||
INSERT INTO functions (name, file_path, body, start_line, end_line, parent_function, parent_class)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
entry.name,
|
||||
entry.file_path,
|
||||
entry.body,
|
||||
entry.start_line,
|
||||
entry.end_line,
|
||||
entry.parent_function,
|
||||
entry.parent_class,
|
||||
),
|
||||
)
|
||||
|
||||
def _insert_class(self, entry: ClassEntry) -> None:
|
||||
"""
|
||||
Insert class entry into db.
|
||||
|
||||
Args:
|
||||
entry: the entry to insert
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
self._db_connection.execute(
|
||||
"""
|
||||
INSERT INTO classes (name, file_path, body, fields, methods, start_line, end_line)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
entry.name,
|
||||
entry.file_path,
|
||||
entry.body,
|
||||
entry.fields,
|
||||
entry.methods,
|
||||
entry.start_line,
|
||||
entry.end_line,
|
||||
),
|
||||
)
|
||||
|
||||
def query_function(
|
||||
self, identifier: str, entry_type: Literal["function", "class_method"] = "function"
|
||||
) -> list[FunctionEntry]:
|
||||
"""
|
||||
Search for a function in the database.
|
||||
|
||||
Args:
|
||||
identifier: the identifier of the function to search for
|
||||
|
||||
Returns:
|
||||
a list of function entries
|
||||
"""
|
||||
records = self._db_connection.execute(
|
||||
"""SELECT name, file_path, body, start_line, end_line, parent_function, parent_class FROM functions WHERE name = ?""",
|
||||
(identifier,),
|
||||
).fetchall()
|
||||
function_entries: list[FunctionEntry] = []
|
||||
for record in records:
|
||||
match entry_type:
|
||||
case "function":
|
||||
if record[6] is None:
|
||||
function_entries.append(
|
||||
FunctionEntry(
|
||||
name=record[0],
|
||||
file_path=record[1],
|
||||
body=record[2],
|
||||
start_line=record[3],
|
||||
end_line=record[4],
|
||||
parent_function=record[5],
|
||||
parent_class=record[6],
|
||||
)
|
||||
)
|
||||
case "class_method":
|
||||
if record[6] is not None:
|
||||
function_entries.append(
|
||||
FunctionEntry(
|
||||
name=record[0],
|
||||
file_path=record[1],
|
||||
body=record[2],
|
||||
start_line=record[3],
|
||||
end_line=record[4],
|
||||
parent_function=record[5],
|
||||
parent_class=record[6],
|
||||
)
|
||||
)
|
||||
return function_entries
|
||||
|
||||
def query_class(self, identifier: str) -> list[ClassEntry]:
|
||||
"""
|
||||
Search for a class in the database.
|
||||
|
||||
Args:
|
||||
identifier: the identifier of the class to search for
|
||||
|
||||
Returns:
|
||||
a list of class entries
|
||||
"""
|
||||
records = self._db_connection.execute(
|
||||
"""SELECT name, file_path, body, fields, methods, start_line, end_line FROM classes WHERE name = ?""",
|
||||
(identifier,),
|
||||
).fetchall()
|
||||
class_entries: list[ClassEntry] = []
|
||||
for record in records:
|
||||
class_entries.append(
|
||||
ClassEntry(
|
||||
name=record[0],
|
||||
file_path=record[1],
|
||||
body=record[2],
|
||||
fields=record[3],
|
||||
methods=record[4],
|
||||
start_line=record[5],
|
||||
end_line=record[6],
|
||||
)
|
||||
)
|
||||
return class_entries
|
||||
@@ -0,0 +1,224 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from pathlib import Path
|
||||
from typing import override
|
||||
|
||||
from trae_agent.tools.base import Tool, ToolCallArguments, ToolExecResult, ToolParameter
|
||||
from trae_agent.tools.ckg.ckg_database import CKGDatabase
|
||||
from trae_agent.tools.run import MAX_RESPONSE_LEN
|
||||
|
||||
CKGToolCommands = ["search_function", "search_class", "search_class_method"]
|
||||
|
||||
|
||||
class CKGTool(Tool):
|
||||
"""Tool to construct and query the code knowledge graph of a codebase."""
|
||||
|
||||
def __init__(self, model_provider: str | None = None) -> None:
|
||||
super().__init__(model_provider)
|
||||
|
||||
# We store the codebase path with built CKG in the following format:
|
||||
# {
|
||||
# "codebase_path": {
|
||||
# "db_connection": sqlite3.Connection,
|
||||
# "codebase_snapshot_hash": str,
|
||||
# }
|
||||
# }
|
||||
self._ckg_databases: dict[Path, CKGDatabase] = {}
|
||||
|
||||
@override
|
||||
def get_model_provider(self) -> str | None:
|
||||
return self._model_provider
|
||||
|
||||
@override
|
||||
def get_name(self) -> str:
|
||||
return "ckg"
|
||||
|
||||
@override
|
||||
def get_description(self) -> str:
|
||||
return """Query the code knowledge graph of a codebase.
|
||||
* State is persistent across command calls and discussions with the user
|
||||
* The `search_function` command searches for functions in the codebase
|
||||
* The `search_class` command searches for classes in the codebase
|
||||
* The `search_class_method` command searches for class methods in the codebase
|
||||
* If a `command` generates a long output, it will be truncated and marked with `<response clipped>`
|
||||
* If multiple entries are found, the tool will return all of them until the truncation is reached.
|
||||
* By default, the tool will print function or class bodies as well as the file path and line number of the function or class. You can disable this by setting the `print_body` parameter to `false`.
|
||||
* The CKG is not completely accurate, and may not be able to find all functions or classes in the codebase.
|
||||
"""
|
||||
|
||||
@override
|
||||
def get_parameters(self) -> list[ToolParameter]:
|
||||
return [
|
||||
ToolParameter(
|
||||
name="command",
|
||||
type="string",
|
||||
description=f"The command to run. Allowed options are {', '.join(CKGToolCommands)}.",
|
||||
required=True,
|
||||
enum=CKGToolCommands,
|
||||
),
|
||||
ToolParameter(
|
||||
name="path",
|
||||
type="string",
|
||||
description="The path to the codebase.",
|
||||
required=True,
|
||||
),
|
||||
ToolParameter(
|
||||
name="identifier",
|
||||
type="string",
|
||||
description="The identifier of the function or class to search for in the code knowledge graph.",
|
||||
required=True,
|
||||
),
|
||||
ToolParameter(
|
||||
name="print_body",
|
||||
type="boolean",
|
||||
description="Whether to print the body of the function or class. This is enabled by default.",
|
||||
required=False,
|
||||
),
|
||||
]
|
||||
|
||||
@override
|
||||
async def execute(self, arguments: ToolCallArguments) -> ToolExecResult:
|
||||
command = str(arguments.get("command")) if "command" in arguments else None
|
||||
if command is None:
|
||||
return ToolExecResult(
|
||||
error=f"No command provided for the {self.get_name()} tool",
|
||||
error_code=-1,
|
||||
)
|
||||
path = str(arguments.get("path")) if "path" in arguments else None
|
||||
if path is None:
|
||||
return ToolExecResult(
|
||||
error=f"No path provided for the {self.get_name()} tool",
|
||||
error_code=-1,
|
||||
)
|
||||
identifier = str(arguments.get("identifier")) if "identifier" in arguments else None
|
||||
if identifier is None:
|
||||
return ToolExecResult(
|
||||
error=f"No identifier provided for the {self.get_name()} tool",
|
||||
error_code=-1,
|
||||
)
|
||||
print_body = bool(arguments.get("print_body")) if "print_body" in arguments else True
|
||||
|
||||
codebase_path = Path(path)
|
||||
if not codebase_path.exists():
|
||||
return ToolExecResult(
|
||||
error=f"Codebase path {path} does not exist",
|
||||
error_code=-1,
|
||||
)
|
||||
if not codebase_path.is_dir():
|
||||
return ToolExecResult(
|
||||
error=f"Codebase path {path} is not a directory",
|
||||
error_code=-1,
|
||||
)
|
||||
|
||||
ckg_database = self._ckg_databases.get(codebase_path)
|
||||
if ckg_database is None:
|
||||
ckg_database = CKGDatabase(codebase_path)
|
||||
self._ckg_databases[codebase_path] = ckg_database
|
||||
|
||||
match command:
|
||||
case "search_function":
|
||||
return ToolExecResult(
|
||||
output=self._search_function(ckg_database, identifier, print_body)
|
||||
)
|
||||
case "search_class":
|
||||
return ToolExecResult(
|
||||
output=self._search_class(ckg_database, identifier, print_body)
|
||||
)
|
||||
case "search_class_method":
|
||||
return ToolExecResult(
|
||||
output=self._search_class_method(ckg_database, identifier, print_body)
|
||||
)
|
||||
case _:
|
||||
return ToolExecResult(error=f"Invalid command: {command}", error_code=-1)
|
||||
|
||||
def _search_function(
|
||||
self, ckg_database: CKGDatabase, identifier: str, print_body: bool = True
|
||||
) -> str:
|
||||
"""Search for a function in the ckg database."""
|
||||
|
||||
entries = ckg_database.query_function(identifier, entry_type="function")
|
||||
|
||||
if len(entries) == 0:
|
||||
return f"No functions named {identifier} found."
|
||||
|
||||
output = f"Found {len(entries)} functions named {identifier}:\n"
|
||||
|
||||
index = 1
|
||||
for entry in entries:
|
||||
output += f"{index}. {entry.file_path}:{entry.start_line}-{entry.end_line}\n"
|
||||
if print_body:
|
||||
output += f"{entry.body}\n\n"
|
||||
|
||||
index += 1
|
||||
|
||||
if len(output) > MAX_RESPONSE_LEN:
|
||||
output = (
|
||||
output[:MAX_RESPONSE_LEN]
|
||||
+ f"\n<response clipped> {len(entries) - index + 1} more entries not shown"
|
||||
)
|
||||
break
|
||||
|
||||
return output
|
||||
|
||||
def _search_class(
|
||||
self, ckg_database: CKGDatabase, identifier: str, print_body: bool = True
|
||||
) -> str:
|
||||
"""Search for a class in the ckg database."""
|
||||
|
||||
entries = ckg_database.query_class(identifier)
|
||||
|
||||
if len(entries) == 0:
|
||||
return f"No classes named {identifier} found."
|
||||
|
||||
output = f"Found {len(entries)} classes named {identifier}:\n"
|
||||
|
||||
index = 1
|
||||
for entry in entries:
|
||||
output += f"{index}. {entry.file_path}:{entry.start_line}-{entry.end_line}\n"
|
||||
if entry.fields:
|
||||
output += f"Fields:\n{entry.fields}\n"
|
||||
if entry.methods:
|
||||
output += f"Methods:\n{entry.methods}\n"
|
||||
if print_body:
|
||||
output += f"{entry.body}\n\n"
|
||||
|
||||
index += 1
|
||||
|
||||
if len(output) > MAX_RESPONSE_LEN:
|
||||
output = (
|
||||
output[:MAX_RESPONSE_LEN]
|
||||
+ f"\n<response clipped> {len(entries) - index + 1} more entries not shown"
|
||||
)
|
||||
break
|
||||
|
||||
return output
|
||||
|
||||
def _search_class_method(
|
||||
self, ckg_database: CKGDatabase, identifier: str, print_body: bool = True
|
||||
) -> str:
|
||||
"""Search for a class method in the ckg database."""
|
||||
|
||||
entries = ckg_database.query_function(identifier, entry_type="class_method")
|
||||
|
||||
if len(entries) == 0:
|
||||
return f"No class methods named {identifier} found."
|
||||
|
||||
output = f"Found {len(entries)} class methods named {identifier}:\n"
|
||||
|
||||
index = 1
|
||||
for entry in entries:
|
||||
output += f"{index}. {entry.file_path}:{entry.start_line}-{entry.end_line} within class {entry.parent_class}\n"
|
||||
if print_body:
|
||||
output += f"{entry.body}\n\n"
|
||||
|
||||
index += 1
|
||||
|
||||
if len(output) > MAX_RESPONSE_LEN:
|
||||
output = (
|
||||
output[:MAX_RESPONSE_LEN]
|
||||
+ f"\n<response clipped> {len(entries) - index + 1} more entries not shown"
|
||||
)
|
||||
break
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,163 @@
|
||||
import json
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from trae_agent.agent.docker_manager import DockerManager
|
||||
from trae_agent.tools.base import ToolCall, ToolExecutor, ToolResult
|
||||
|
||||
|
||||
class DockerToolExecutor:
|
||||
"""
|
||||
A ToolExecutor that delegates tool calls to either a local executor
|
||||
or a Docker environment based on the tool's name.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
original_executor: ToolExecutor,
|
||||
docker_manager: DockerManager,
|
||||
docker_tools: list[str],
|
||||
host_workspace_dir: str | None,
|
||||
container_workspace_dir: str,
|
||||
):
|
||||
"""
|
||||
Initializes the DockerToolExecutor.
|
||||
"""
|
||||
self._original_executor = original_executor
|
||||
self._docker_manager = docker_manager
|
||||
self._docker_tools_set = set(docker_tools)
|
||||
# Get path from __init__ ---
|
||||
self._host_workspace_dir = (
|
||||
os.path.abspath(host_workspace_dir) if host_workspace_dir else None
|
||||
)
|
||||
self._container_workspace_dir = container_workspace_dir
|
||||
|
||||
def _translate_path(self, host_path: str) -> str:
|
||||
"""Robust path translation function: Translate the host path into the corresponding path within the container."""
|
||||
if not self._host_workspace_dir:
|
||||
return host_path # 如果没有配置主机工作区，则不翻译
|
||||
abs_host_path = os.path.abspath(host_path)
|
||||
if (
|
||||
os.path.commonpath([abs_host_path, self._host_workspace_dir])
|
||||
== self._host_workspace_dir
|
||||
):
|
||||
relative_path = os.path.relpath(abs_host_path, self._host_workspace_dir)
|
||||
container_path = os.path.join(self._container_workspace_dir, relative_path)
|
||||
return os.path.normpath(container_path)
|
||||
return host_path
|
||||
|
||||
async def close_tools(self):
|
||||
"""
|
||||
Closes any resources held by the underlying original executor.
|
||||
This method fulfills the contract expected by BaseAgent.
|
||||
"""
|
||||
if self._original_executor:
|
||||
return await self._original_executor.close_tools()
|
||||
|
||||
async def sequential_tool_call(self, tool_calls: list[ToolCall]) -> list[ToolResult]:
|
||||
"""Executes tool calls sequentially, routing to Docker if necessary."""
|
||||
results = []
|
||||
for tool_call in tool_calls:
|
||||
if tool_call.name in self._docker_tools_set:
|
||||
result = self._execute_in_docker(tool_call)
|
||||
else:
|
||||
# Execute locally
|
||||
result_list = await self._original_executor.sequential_tool_call([tool_call])
|
||||
result = result_list[0]
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
async def parallel_tool_call(self, tool_calls: list[ToolCall]) -> list[ToolResult]:
|
||||
"""For simplicity, parallel calls are also executed sequentially."""
|
||||
# print(
|
||||
# "[yellow]Warning: Parallel tool calls are executed sequentially in Docker mode.[/yellow]"
|
||||
# )
|
||||
return await self.sequential_tool_call(tool_calls)
|
||||
|
||||
def _execute_in_docker(self, tool_call: ToolCall) -> ToolResult:
|
||||
"""
|
||||
Builds and executes a command inside the Docker container,
|
||||
with path translation.
|
||||
"""
|
||||
try:
|
||||
# --- Parameter preprocessing and path translation ---
|
||||
processed_args: dict[str, Any] = {}
|
||||
for key, value in tool_call.arguments.items():
|
||||
# Assuming that all parameters named 'path' are paths that need to be translated
|
||||
if key == "path" and isinstance(value, str):
|
||||
translated_path = self._translate_path(value)
|
||||
processed_args[key] = translated_path
|
||||
else:
|
||||
processed_args[key] = value
|
||||
|
||||
# --- The subsequent logic now uses' processed'args' instead of 'tool_call. arguments' ---
|
||||
command_to_run = ""
|
||||
|
||||
# --- Rule 1: Handling bash tools ---
|
||||
if tool_call.name == "bash":
|
||||
command_value = processed_args.get("command")
|
||||
if not isinstance(command_value, str) or not command_value:
|
||||
raise ValueError("Tool 'bash' requires a non-empty 'command' string argument.")
|
||||
command_to_run = command_value
|
||||
|
||||
# --- Rule2 : Handling str_replace_based_edit_tool ---
|
||||
elif tool_call.name == "str_replace_based_edit_tool":
|
||||
sub_command = processed_args.get("command")
|
||||
if not sub_command:
|
||||
raise ValueError("Edit tool called without a 'command' (sub-command).")
|
||||
|
||||
if not isinstance(sub_command, str):
|
||||
raise TypeError(
|
||||
f"The 'command' argument for {tool_call.name} must be a string."
|
||||
)
|
||||
executable_path = f"{self._docker_manager.CONTAINER_TOOLS_PATH}/edit_tool"
|
||||
cmd_parts = [executable_path, sub_command]
|
||||
|
||||
for key, value in processed_args.items():
|
||||
if key == "command" or value is None:
|
||||
continue
|
||||
if isinstance(value, list):
|
||||
str_value = " ".join(map(str, value))
|
||||
cmd_parts.append(f"--{key} {str_value}")
|
||||
else:
|
||||
cmd_parts.append(f"--{key} '{str(value)}'")
|
||||
|
||||
command_to_run = " ".join(cmd_parts)
|
||||
# --- Rule 3: Handling json_edit_tool ---
|
||||
elif tool_call.name == "json_edit_tool":
|
||||
executable_path = f"{self._docker_manager.CONTAINER_TOOLS_PATH}/json_edit_tool"
|
||||
cmd_parts = [executable_path]
|
||||
for key, value in processed_args.items():
|
||||
if value is None:
|
||||
continue
|
||||
# --- Serialize the 'value' parameter into a JSON string ---
|
||||
if key == "value":
|
||||
json_string_value = json.dumps(value)
|
||||
cmd_parts.append(f"--{key} '{json_string_value}'")
|
||||
elif isinstance(value, list):
|
||||
# In theory, json edit_tool does not have a list parameter, but it should be kept as a precautionary measure
|
||||
cmd_parts.append(f"--{key} {' '.join(map(str, value))}")
|
||||
else:
|
||||
cmd_parts.append(f"--{key} '{str(value)}'")
|
||||
command_to_run = " ".join(cmd_parts)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"The logic for Docker execution of tool '{tool_call.name}' is not implemented."
|
||||
)
|
||||
|
||||
# Execute the final built command
|
||||
exit_code, output = self._docker_manager.execute(command_to_run)
|
||||
return ToolResult(
|
||||
call_id=tool_call.call_id,
|
||||
name=tool_call.name,
|
||||
result=output,
|
||||
success=exit_code == 0,
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolResult(
|
||||
call_id=tool_call.call_id,
|
||||
name=tool_call.name,
|
||||
result=f"Failed to build or execute command for tool '{tool_call.name}' in Docker: {e}",
|
||||
success=False,
|
||||
error=str(e),
|
||||
)
|
||||
@@ -0,0 +1,360 @@
|
||||
# Copyright (c) 2023 Anthropic
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: MIT
|
||||
#
|
||||
# This file has been modified by ByteDance Ltd. and/or its affiliates. on 13 June 2025
|
||||
#
|
||||
# Original file was released under MIT License, with the full license text
|
||||
# available at https://github.com/anthropics/anthropic-quickstarts/blob/main/LICENSE
|
||||
#
|
||||
# This modified file is released under the same license.
|
||||
|
||||
from pathlib import Path
|
||||
from typing import override
|
||||
|
||||
from trae_agent.tools.base import Tool, ToolCallArguments, ToolError, ToolExecResult, ToolParameter
|
||||
from trae_agent.tools.run import maybe_truncate, run
|
||||
|
||||
EditToolSubCommands = [
|
||||
"view",
|
||||
"create",
|
||||
"str_replace",
|
||||
"insert",
|
||||
]
|
||||
SNIPPET_LINES: int = 4
|
||||
|
||||
|
||||
class TextEditorTool(Tool):
|
||||
"""Tool to replace a string in a file."""
|
||||
|
||||
def __init__(self, model_provider: str | None = None) -> None:
|
||||
super().__init__(model_provider)
|
||||
|
||||
@override
|
||||
def get_model_provider(self) -> str | None:
|
||||
return self._model_provider
|
||||
|
||||
@override
|
||||
def get_name(self) -> str:
|
||||
return "str_replace_based_edit_tool"
|
||||
|
||||
@override
|
||||
def get_description(self) -> str:
|
||||
return """Custom editing tool for viewing, creating and editing files
|
||||
* State is persistent across command calls and discussions with the user
|
||||
* If `path` is a file, `view` displays the result of applying `cat -n`. If `path` is a directory, `view` lists non-hidden files and directories up to 2 levels deep
|
||||
* The `create` command cannot be used if the specified `path` already exists as a file !!! If you know that the `path` already exists, please remove it first and then perform the `create` operation!
|
||||
* If a `command` generates a long output, it will be truncated and marked with `<response clipped>`
|
||||
|
||||
Notes for using the `str_replace` command:
|
||||
* The `old_str` parameter should match EXACTLY one or more consecutive lines from the original file. Be mindful of whitespaces!
|
||||
* If the `old_str` parameter is not unique in the file, the replacement will not be performed. Make sure to include enough context in `old_str` to make it unique
|
||||
* The `new_str` parameter should contain the edited lines that should replace the `old_str`
|
||||
"""
|
||||
|
||||
@override
|
||||
def get_parameters(self) -> list[ToolParameter]:
|
||||
"""Get the parameters for the str_replace_based_edit_tool."""
|
||||
return [
|
||||
ToolParameter(
|
||||
name="command",
|
||||
type="string",
|
||||
description=f"The commands to run. Allowed options are: {', '.join(EditToolSubCommands)}.",
|
||||
required=True,
|
||||
enum=EditToolSubCommands,
|
||||
),
|
||||
ToolParameter(
|
||||
name="file_text",
|
||||
type="string",
|
||||
description="Required parameter of `create` command, with the content of the file to be created.",
|
||||
),
|
||||
ToolParameter(
|
||||
name="insert_line",
|
||||
type="integer",
|
||||
description="Required parameter of `insert` command. The `new_str` will be inserted AFTER the line `insert_line` of `path`.",
|
||||
),
|
||||
ToolParameter(
|
||||
name="new_str",
|
||||
type="string",
|
||||
description="Optional parameter of `str_replace` command containing the new string (if not given, no string will be added). Required parameter of `insert` command containing the string to insert.",
|
||||
),
|
||||
ToolParameter(
|
||||
name="old_str",
|
||||
type="string",
|
||||
description="Required parameter of `str_replace` command containing the string in `path` to replace.",
|
||||
),
|
||||
ToolParameter(
|
||||
name="path",
|
||||
type="string",
|
||||
description="Absolute path to file or directory, e.g. `/repo/file.py` or `/repo`.",
|
||||
required=True,
|
||||
),
|
||||
ToolParameter(
|
||||
name="view_range",
|
||||
type="array",
|
||||
description="Optional parameter of `view` command when `path` points to a file. If none is given, the full file is shown. If provided, the file will be shown in the indicated line number range, e.g. [11, 12] will show lines 11 and 12. Indexing at 1 to start. Setting `[start_line, -1]` shows all lines from `start_line` to the end of the file.",
|
||||
items={"type": "integer"},
|
||||
),
|
||||
]
|
||||
|
||||
@override
|
||||
async def execute(self, arguments: ToolCallArguments) -> ToolExecResult:
|
||||
"""Execute the str_replace_editor tool."""
|
||||
command = str(arguments["command"]) if "command" in arguments else None
|
||||
if command is None:
|
||||
return ToolExecResult(
|
||||
error=f"No command provided for the {self.get_name()} tool",
|
||||
error_code=-1,
|
||||
)
|
||||
path = str(arguments["path"]) if "path" in arguments else None
|
||||
if path is None:
|
||||
return ToolExecResult(
|
||||
error=f"No path provided for the {self.get_name()} tool", error_code=-1
|
||||
)
|
||||
_path = Path(path)
|
||||
try:
|
||||
self.validate_path(command, _path)
|
||||
match command:
|
||||
case "view":
|
||||
return await self._view_handler(arguments, _path)
|
||||
case "create":
|
||||
return self._create_handler(arguments, _path)
|
||||
case "str_replace":
|
||||
return self._str_replace_handler(arguments, _path)
|
||||
case "insert":
|
||||
return self._insert_handler(arguments, _path)
|
||||
case _:
|
||||
return ToolExecResult(
|
||||
error=f"Unrecognized command {command}. The allowed commands for the {self.name} tool are: {', '.join(EditToolSubCommands)}",
|
||||
error_code=-1,
|
||||
)
|
||||
except ToolError as e:
|
||||
return ToolExecResult(error=str(e), error_code=-1)
|
||||
|
||||
def validate_path(self, command: str, path: Path):
|
||||
"""Validate the path for the str_replace_editor tool."""
|
||||
if not path.is_absolute():
|
||||
suggested_path = Path("/") / path
|
||||
raise ToolError(
|
||||
f"The path {path} is not an absolute path, it should start with `/`. Maybe you meant {suggested_path}?"
|
||||
)
|
||||
# Check if path exists
|
||||
if not path.exists() and command != "create":
|
||||
raise ToolError(f"The path {path} does not exist. Please provide a valid path.")
|
||||
if path.exists() and command == "create":
|
||||
raise ToolError(
|
||||
f"File already exists at: {path}. Cannot overwrite files using command `create`."
|
||||
)
|
||||
# Check if the path points to a directory
|
||||
if path.is_dir() and command != "view":
|
||||
raise ToolError(
|
||||
f"The path {path} is a directory and only the `view` command can be used on directories"
|
||||
)
|
||||
|
||||
async def _view(self, path: Path, view_range: list[int] | None = None) -> ToolExecResult:
|
||||
"""Implement the view command"""
|
||||
if path.is_dir():
|
||||
if view_range:
|
||||
raise ToolError(
|
||||
"The `view_range` parameter is not allowed when `path` points to a directory."
|
||||
)
|
||||
|
||||
return_code, stdout, stderr = await run(rf"find {path} -maxdepth 2 -not -path '*/\.*'")
|
||||
if not stderr:
|
||||
stdout = f"Here's the files and directories up to 2 levels deep in {path}, excluding hidden items:\n{stdout}\n"
|
||||
return ToolExecResult(error_code=return_code, output=stdout, error=stderr)
|
||||
|
||||
file_content = self.read_file(path)
|
||||
init_line = 1
|
||||
if view_range:
|
||||
if len(view_range) != 2 or not all(isinstance(i, int) for i in view_range): # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
raise ToolError("Invalid `view_range`. It should be a list of two integers.")
|
||||
file_lines = file_content.split("\n")
|
||||
n_lines_file = len(file_lines)
|
||||
init_line, final_line = view_range
|
||||
if init_line < 1 or init_line > n_lines_file:
|
||||
raise ToolError(
|
||||
f"Invalid `view_range`: {view_range}. Its first element `{init_line}` should be within the range of lines of the file: {[1, n_lines_file]}"
|
||||
)
|
||||
if final_line > n_lines_file:
|
||||
raise ToolError(
|
||||
f"Invalid `view_range`: {view_range}. Its second element `{final_line}` should be smaller than the number of lines in the file: `{n_lines_file}`"
|
||||
)
|
||||
if final_line != -1 and final_line < init_line:
|
||||
raise ToolError(
|
||||
f"Invalid `view_range`: {view_range}. Its second element `{final_line}` should be larger or equal than its first `{init_line}`"
|
||||
)
|
||||
|
||||
if final_line == -1:
|
||||
file_content = "\n".join(file_lines[init_line - 1 :])
|
||||
else:
|
||||
file_content = "\n".join(file_lines[init_line - 1 : final_line])
|
||||
|
||||
return ToolExecResult(
|
||||
output=self._make_output(file_content, str(path), init_line=init_line)
|
||||
)
|
||||
|
||||
def str_replace(self, path: Path, old_str: str, new_str: str | None) -> ToolExecResult:
|
||||
"""Implement the str_replace command, which replaces old_str with new_str in the file content"""
|
||||
# Read the file content
|
||||
file_content = self.read_file(path).expandtabs()
|
||||
old_str = old_str.expandtabs()
|
||||
new_str = new_str.expandtabs() if new_str is not None else ""
|
||||
|
||||
# Check if old_str is unique in the file
|
||||
occurrences = file_content.count(old_str)
|
||||
if occurrences == 0:
|
||||
raise ToolError(
|
||||
f"No replacement was performed, old_str `{old_str}` did not appear verbatim in {path}."
|
||||
)
|
||||
elif occurrences > 1:
|
||||
file_content_lines = file_content.split("\n")
|
||||
lines = [idx + 1 for idx, line in enumerate(file_content_lines) if old_str in line]
|
||||
raise ToolError(
|
||||
f"No replacement was performed. Multiple occurrences of old_str `{old_str}` in lines {lines}. Please ensure it is unique"
|
||||
)
|
||||
|
||||
# Replace old_str with new_str
|
||||
new_file_content = file_content.replace(old_str, new_str)
|
||||
|
||||
# Write the new content to the file
|
||||
self.write_file(path, new_file_content)
|
||||
|
||||
# Create a snippet of the edited section
|
||||
replacement_line = file_content.split(old_str)[0].count("\n")
|
||||
start_line = max(0, replacement_line - SNIPPET_LINES)
|
||||
end_line = replacement_line + SNIPPET_LINES + new_str.count("\n")
|
||||
snippet = "\n".join(new_file_content.split("\n")[start_line : end_line + 1])
|
||||
|
||||
# Prepare the success message
|
||||
success_msg = f"The file {path} has been edited. "
|
||||
success_msg += self._make_output(snippet, f"a snippet of {path}", start_line + 1)
|
||||
success_msg += "Review the changes and make sure they are as expected. Edit the file again if necessary."
|
||||
|
||||
return ToolExecResult(
|
||||
output=success_msg,
|
||||
)
|
||||
|
||||
def _insert(self, path: Path, insert_line: int, new_str: str) -> ToolExecResult:
|
||||
"""Implement the insert command, which inserts new_str at the specified line in the file content."""
|
||||
file_text = self.read_file(path).expandtabs()
|
||||
new_str = new_str.expandtabs()
|
||||
file_text_lines = file_text.split("\n")
|
||||
n_lines_file = len(file_text_lines)
|
||||
|
||||
if insert_line < 0 or insert_line > n_lines_file:
|
||||
raise ToolError(
|
||||
f"Invalid `insert_line` parameter: {insert_line}. It should be within the range of lines of the file: {[0, n_lines_file]}"
|
||||
)
|
||||
|
||||
new_str_lines = new_str.split("\n")
|
||||
new_file_text_lines = (
|
||||
file_text_lines[:insert_line] + new_str_lines + file_text_lines[insert_line:]
|
||||
)
|
||||
snippet_lines = (
|
||||
file_text_lines[max(0, insert_line - SNIPPET_LINES) : insert_line]
|
||||
+ new_str_lines
|
||||
+ file_text_lines[insert_line : insert_line + SNIPPET_LINES]
|
||||
)
|
||||
|
||||
new_file_text = "\n".join(new_file_text_lines)
|
||||
snippet = "\n".join(snippet_lines)
|
||||
|
||||
self.write_file(path, new_file_text)
|
||||
|
||||
success_msg = f"The file {path} has been edited. "
|
||||
success_msg += self._make_output(
|
||||
snippet,
|
||||
"a snippet of the edited file",
|
||||
max(1, insert_line - SNIPPET_LINES + 1),
|
||||
)
|
||||
success_msg += "Review the changes and make sure they are as expected (correct indentation, no duplicate lines, etc). Edit the file again if necessary."
|
||||
return ToolExecResult(
|
||||
output=success_msg,
|
||||
)
|
||||
|
||||
# Note: undo_edit method is not implemented in this version as it was removed
|
||||
|
||||
def read_file(self, path: Path):
|
||||
"""Read the content of a file from a given path; raise a ToolError if an error occurs."""
|
||||
try:
|
||||
return path.read_text()
|
||||
except Exception as e:
|
||||
raise ToolError(f"Ran into {e} while trying to read {path}") from None
|
||||
|
||||
def write_file(self, path: Path, file: str):
|
||||
"""Write the content of a file to a given path; raise a ToolError if an error occurs."""
|
||||
try:
|
||||
_ = path.write_text(file)
|
||||
except Exception as e:
|
||||
raise ToolError(f"Ran into {e} while trying to write to {path}") from None
|
||||
|
||||
def _make_output(
|
||||
self,
|
||||
file_content: str,
|
||||
file_descriptor: str,
|
||||
init_line: int = 1,
|
||||
expand_tabs: bool = True,
|
||||
):
|
||||
"""Generate output for the CLI based on the content of a file."""
|
||||
file_content = maybe_truncate(file_content)
|
||||
if expand_tabs:
|
||||
file_content = file_content.expandtabs()
|
||||
file_content = "\n".join(
|
||||
[f"{i + init_line:6}\t{line}" for i, line in enumerate(file_content.split("\n"))]
|
||||
)
|
||||
return (
|
||||
f"Here's the result of running `cat -n` on {file_descriptor}:\n" + file_content + "\n"
|
||||
)
|
||||
|
||||
async def _view_handler(self, arguments: ToolCallArguments, _path: Path) -> ToolExecResult:
|
||||
view_range = arguments.get("view_range", None)
|
||||
if view_range is None:
|
||||
return await self._view(_path, None)
|
||||
if not (isinstance(view_range, list) and all(isinstance(i, int) for i in view_range)):
|
||||
return ToolExecResult(
|
||||
error="Parameter `view_range` should be a list of integers.",
|
||||
error_code=-1,
|
||||
)
|
||||
view_range_int: list[int] = [i for i in view_range if isinstance(i, int)]
|
||||
return await self._view(_path, view_range_int)
|
||||
|
||||
def _create_handler(self, arguments: ToolCallArguments, _path: Path) -> ToolExecResult:
|
||||
file_text = arguments.get("file_text", None)
|
||||
if not isinstance(file_text, str):
|
||||
return ToolExecResult(
|
||||
error="Parameter `file_text` is required and must be a string for command: create",
|
||||
error_code=-1,
|
||||
)
|
||||
self.write_file(_path, file_text)
|
||||
return ToolExecResult(output=f"File created successfully at: {_path}")
|
||||
|
||||
def _str_replace_handler(self, arguments: ToolCallArguments, _path: Path) -> ToolExecResult:
|
||||
old_str = arguments.get("old_str") if "old_str" in arguments else None
|
||||
if not isinstance(old_str, str):
|
||||
return ToolExecResult(
|
||||
error="Parameter `old_str` is required and should be a string for command: str_replace",
|
||||
error_code=-1,
|
||||
)
|
||||
new_str = arguments.get("new_str") if "new_str" in arguments else None
|
||||
if not (new_str is None or isinstance(new_str, str)):
|
||||
return ToolExecResult(
|
||||
error="Parameter `new_str` should be a string or null for command: str_replace",
|
||||
error_code=-1,
|
||||
)
|
||||
return self.str_replace(_path, old_str, new_str)
|
||||
|
||||
def _insert_handler(self, arguments: ToolCallArguments, _path: Path) -> ToolExecResult:
|
||||
insert_line = arguments.get("insert_line") if "insert_line" in arguments else None
|
||||
if not isinstance(insert_line, int):
|
||||
return ToolExecResult(
|
||||
error="Parameter `insert_line` is required and should be integer for command: insert",
|
||||
error_code=-1,
|
||||
)
|
||||
new_str_to_insert = arguments.get("new_str") if "new_str" in arguments else None
|
||||
if not isinstance(new_str_to_insert, str):
|
||||
return ToolExecResult(
|
||||
error="Parameter `new_str` is required for command: insert",
|
||||
error_code=-1,
|
||||
)
|
||||
return self._insert(_path, insert_line, new_str_to_insert)
|
||||
@@ -0,0 +1,477 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
# Dependency Definition Area: Here we define all the required "blueprints" and "parts"
|
||||
# This is a minimal 'override' alternative. Since we no longer need it after packaging, we can define a function that does nothing
|
||||
def override(f):
|
||||
return f
|
||||
|
||||
|
||||
# A simple base class that makes' class TextEditorTool (Tool): 'grammatically correct
|
||||
class Tool:
|
||||
def __init__(self, model_provider: str | None = None) -> None:
|
||||
self._model_provider = model_provider
|
||||
|
||||
|
||||
# ToolCallArguments is just a type alias, we can use dict instead
|
||||
ToolCallArguments = dict
|
||||
|
||||
|
||||
# Custom exception class
|
||||
class ToolError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
# A class used to encapsulate the results of tool execution
|
||||
class ToolExecResult:
|
||||
def __init__(self, output: str | None = None, error: str | None = None, error_code: int = 0):
|
||||
self.output = output
|
||||
self.error = error
|
||||
self.error_code = error_code
|
||||
|
||||
|
||||
# Class used to describe tool parameters (although not directly used in CLI, TextEditTool's methods require it)
|
||||
class ToolParameter:
|
||||
def __init__(self, name: str, type: str, description: str, required: bool = False, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
def maybe_truncate(output: str, max_chars: int = 20000) -> str:
|
||||
"""Truncate the output if it's too long."""
|
||||
if len(output) > max_chars:
|
||||
return output[:max_chars] + "\n<... response clipped ...>\n"
|
||||
return output
|
||||
|
||||
|
||||
EditToolSubCommands = ["view", "create", "str_replace", "insert"]
|
||||
SNIPPET_LINES = 5
|
||||
|
||||
|
||||
async def run(command: str, timeout: int = 300) -> tuple[int, str, str]:
|
||||
"""Run a shell command asynchronously."""
|
||||
proc = await asyncio.create_subprocess_shell(
|
||||
command,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
try:
|
||||
stdout_bytes, stderr_bytes = await asyncio.wait_for(proc.communicate(), timeout=timeout)
|
||||
stdout = stdout_bytes.decode("utf-8", errors="ignore")
|
||||
stderr = stderr_bytes.decode("utf-8", errors="ignore")
|
||||
return proc.returncode if proc.returncode is not None else -1, stdout, stderr
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
return -1, "", f"Command timed out after {timeout} seconds."
|
||||
|
||||
|
||||
class TextEditorTool(Tool):
|
||||
"""Tool to replace a string in a file."""
|
||||
|
||||
def __init__(self, model_provider: str | None = None) -> None:
|
||||
super().__init__(model_provider)
|
||||
|
||||
@override
|
||||
def get_model_provider(self) -> str | None:
|
||||
return self._model_provider
|
||||
|
||||
@override
|
||||
def get_name(self) -> str:
|
||||
return "str_replace_based_edit_tool"
|
||||
|
||||
@override
|
||||
def get_description(self) -> str:
|
||||
return """Custom editing tool for viewing, creating and editing files
|
||||
* State is persistent across command calls and discussions with the user
|
||||
* If `path` is a file, `view` displays the result of applying `cat -n`. If `path` is a directory, `view` lists non-hidden files and directories up to 2 levels deep
|
||||
* The `create` command cannot be used if the specified `path` already exists as a file !!! If you know that the `path` already exists, please remove it first and then perform the `create` operation!
|
||||
* If a `command` generates a long output, it will be truncated and marked with `<response clipped>`
|
||||
|
||||
Notes for using the `str_replace` command:
|
||||
* The `old_str` parameter should match EXACTLY one or more consecutive lines from the original file. Be mindful of whitespaces!
|
||||
* If the `old_str` parameter is not unique in the file, the replacement will not be performed. Make sure to include enough context in `old_str` to make it unique
|
||||
* The `new_str` parameter should contain the edited lines that should replace the `old_str`
|
||||
"""
|
||||
|
||||
@override
|
||||
def get_parameters(self) -> list[ToolParameter]:
|
||||
"""Get the parameters for the str_replace_based_edit_tool."""
|
||||
return [
|
||||
ToolParameter(
|
||||
name="command",
|
||||
type="string",
|
||||
description=f"The commands to run. Allowed options are: {', '.join(EditToolSubCommands)}.",
|
||||
required=True,
|
||||
enum=EditToolSubCommands,
|
||||
),
|
||||
ToolParameter(
|
||||
name="file_text",
|
||||
type="string",
|
||||
description="Required parameter of `create` command, with the content of the file to be created.",
|
||||
),
|
||||
ToolParameter(
|
||||
name="insert_line",
|
||||
type="integer",
|
||||
description="Required parameter of `insert` command. The `new_str` will be inserted AFTER the line `insert_line` of `path`.",
|
||||
),
|
||||
ToolParameter(
|
||||
name="new_str",
|
||||
type="string",
|
||||
description="Optional parameter of `str_replace` command containing the new string (if not given, no string will be added). Required parameter of `insert` command containing the string to insert.",
|
||||
),
|
||||
ToolParameter(
|
||||
name="old_str",
|
||||
type="string",
|
||||
description="Required parameter of `str_replace` command containing the string in `path` to replace.",
|
||||
),
|
||||
ToolParameter(
|
||||
name="path",
|
||||
type="string",
|
||||
description="Absolute path to file or directory, e.g. `/repo/file.py` or `/repo`.",
|
||||
required=True,
|
||||
),
|
||||
ToolParameter(
|
||||
name="view_range",
|
||||
type="array",
|
||||
description="Optional parameter of `view` command when `path` points to a file. If none is given, the full file is shown. If provided, the file will be shown in the indicated line number range, e.g. [11, 12] will show lines 11 and 12. Indexing at 1 to start. Setting `[start_line, -1]` shows all lines from `start_line` to the end of the file.",
|
||||
items={"type": "integer"},
|
||||
),
|
||||
]
|
||||
|
||||
@override
|
||||
async def execute(self, arguments: ToolCallArguments) -> ToolExecResult:
|
||||
"""Execute the str_replace_editor tool."""
|
||||
command = str(arguments["command"]) if "command" in arguments else None
|
||||
if command is None:
|
||||
return ToolExecResult(
|
||||
error=f"No command provided for the {self.get_name()} tool",
|
||||
error_code=-1,
|
||||
)
|
||||
path = str(arguments["path"]) if "path" in arguments else None
|
||||
if path is None:
|
||||
return ToolExecResult(
|
||||
error=f"No path provided for the {self.get_name()} tool", error_code=-1
|
||||
)
|
||||
_path = Path(path)
|
||||
try:
|
||||
self.validate_path(command, _path)
|
||||
match command:
|
||||
case "view":
|
||||
return await self._view_handler(arguments, _path)
|
||||
case "create":
|
||||
return self._create_handler(arguments, _path)
|
||||
case "str_replace":
|
||||
return self._str_replace_handler(arguments, _path)
|
||||
case "insert":
|
||||
return self._insert_handler(arguments, _path)
|
||||
case _:
|
||||
return ToolExecResult(
|
||||
error=f"Unrecognized command {command}. The allowed commands for the {self.get_name()} tool are: {', '.join(EditToolSubCommands)}",
|
||||
error_code=-1,
|
||||
)
|
||||
except ToolError as e:
|
||||
return ToolExecResult(error=str(e), error_code=-1)
|
||||
|
||||
def validate_path(self, command: str, path: Path):
|
||||
"""Validate the path for the str_replace_editor tool."""
|
||||
if not path.is_absolute():
|
||||
suggested_path = Path("/") / path
|
||||
raise ToolError(
|
||||
f"The path {path} is not an absolute path, it should start with `/`. Maybe you meant {suggested_path}?"
|
||||
)
|
||||
# Check if path exists
|
||||
if not path.exists() and command != "create":
|
||||
raise ToolError(f"The path {path} does not exist. Please provide a valid path.")
|
||||
if path.exists() and command == "create":
|
||||
raise ToolError(
|
||||
f"File already exists at: {path}. Cannot overwrite files using command `create`."
|
||||
)
|
||||
# Check if the path points to a directory
|
||||
if path.is_dir() and command != "view":
|
||||
raise ToolError(
|
||||
f"The path {path} is a directory and only the `view` command can be used on directories"
|
||||
)
|
||||
|
||||
async def _view(self, path: Path, view_range: list[int] | None = None) -> ToolExecResult:
|
||||
"""Implement the view command"""
|
||||
if path.is_dir():
|
||||
if view_range:
|
||||
raise ToolError(
|
||||
"The `view_range` parameter is not allowed when `path` points to a directory."
|
||||
)
|
||||
|
||||
return_code, stdout, stderr = await run(rf"find {path} -maxdepth 2 -not -path '*/\.*'")
|
||||
if not stderr:
|
||||
stdout = f"Here's the files and directories up to 2 levels deep in {path}, excluding hidden items:\n{stdout}\n"
|
||||
return ToolExecResult(error_code=return_code, output=stdout, error=stderr)
|
||||
|
||||
file_content = self.read_file(path)
|
||||
init_line = 1
|
||||
if view_range:
|
||||
if len(view_range) != 2 or not all(isinstance(i, int) for i in view_range): # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
raise ToolError("Invalid `view_range`. It should be a list of two integers.")
|
||||
file_lines = file_content.split("\n")
|
||||
n_lines_file = len(file_lines)
|
||||
init_line, final_line = view_range
|
||||
if init_line < 1 or init_line > n_lines_file:
|
||||
raise ToolError(
|
||||
f"Invalid `view_range`: {view_range}. Its first element `{init_line}` should be within the range of lines of the file: {[1, n_lines_file]}"
|
||||
)
|
||||
if final_line > n_lines_file:
|
||||
raise ToolError(
|
||||
f"Invalid `view_range`: {view_range}. Its second element `{final_line}` should be smaller than the number of lines in the file: `{n_lines_file}`"
|
||||
)
|
||||
if final_line != -1 and final_line < init_line:
|
||||
raise ToolError(
|
||||
f"Invalid `view_range`: {view_range}. Its second element `{final_line}` should be larger or equal than its first `{init_line}`"
|
||||
)
|
||||
|
||||
if final_line == -1:
|
||||
file_content = "\n".join(file_lines[init_line - 1 :])
|
||||
else:
|
||||
file_content = "\n".join(file_lines[init_line - 1 : final_line])
|
||||
|
||||
return ToolExecResult(
|
||||
output=self._make_output(file_content, str(path), init_line=init_line)
|
||||
)
|
||||
|
||||
def str_replace(self, path: Path, old_str: str, new_str: str | None) -> ToolExecResult:
|
||||
"""Implement the str_replace command, which replaces old_str with new_str in the file content"""
|
||||
# Read the file content
|
||||
file_content = self.read_file(path).expandtabs()
|
||||
old_str = old_str.expandtabs()
|
||||
new_str = new_str.expandtabs() if new_str is not None else ""
|
||||
|
||||
# Check if old_str is unique in the file
|
||||
occurrences = file_content.count(old_str)
|
||||
if occurrences == 0:
|
||||
raise ToolError(
|
||||
f"No replacement was performed, old_str `{old_str}` did not appear verbatim in {path}."
|
||||
)
|
||||
elif occurrences > 1:
|
||||
file_content_lines = file_content.split("\n")
|
||||
lines = [idx + 1 for idx, line in enumerate(file_content_lines) if old_str in line]
|
||||
raise ToolError(
|
||||
f"No replacement was performed. Multiple occurrences of old_str `{old_str}` in lines {lines}. Please ensure it is unique"
|
||||
)
|
||||
|
||||
# Replace old_str with new_str
|
||||
new_file_content = file_content.replace(old_str, new_str)
|
||||
|
||||
# Write the new content to the file
|
||||
self.write_file(path, new_file_content)
|
||||
|
||||
# Create a snippet of the edited section
|
||||
replacement_line = file_content.split(old_str)[0].count("\n")
|
||||
start_line = max(0, replacement_line - SNIPPET_LINES)
|
||||
end_line = replacement_line + SNIPPET_LINES + new_str.count("\n")
|
||||
snippet = "\n".join(new_file_content.split("\n")[start_line : end_line + 1])
|
||||
|
||||
# Prepare the success message
|
||||
success_msg = f"The file {path} has been edited. "
|
||||
success_msg += self._make_output(snippet, f"a snippet of {path}", start_line + 1)
|
||||
success_msg += "Review the changes and make sure they are as expected. Edit the file again if necessary."
|
||||
|
||||
return ToolExecResult(
|
||||
output=success_msg,
|
||||
)
|
||||
|
||||
def _insert(self, path: Path, insert_line: int, new_str: str) -> ToolExecResult:
|
||||
"""Implement the insert command, which inserts new_str at the specified line in the file content."""
|
||||
file_text = self.read_file(path).expandtabs()
|
||||
new_str = new_str.expandtabs()
|
||||
file_text_lines = file_text.split("\n")
|
||||
n_lines_file = len(file_text_lines)
|
||||
|
||||
if insert_line < 0 or insert_line > n_lines_file:
|
||||
raise ToolError(
|
||||
f"Invalid `insert_line` parameter: {insert_line}. It should be within the range of lines of the file: {[0, n_lines_file]}"
|
||||
)
|
||||
|
||||
new_str_lines = new_str.split("\n")
|
||||
new_file_text_lines = (
|
||||
file_text_lines[:insert_line] + new_str_lines + file_text_lines[insert_line:]
|
||||
)
|
||||
snippet_lines = (
|
||||
file_text_lines[max(0, insert_line - SNIPPET_LINES) : insert_line]
|
||||
+ new_str_lines
|
||||
+ file_text_lines[insert_line : insert_line + SNIPPET_LINES]
|
||||
)
|
||||
|
||||
new_file_text = "\n".join(new_file_text_lines)
|
||||
snippet = "\n".join(snippet_lines)
|
||||
|
||||
self.write_file(path, new_file_text)
|
||||
|
||||
success_msg = f"The file {path} has been edited. "
|
||||
success_msg += self._make_output(
|
||||
snippet,
|
||||
"a snippet of the edited file",
|
||||
max(1, insert_line - SNIPPET_LINES + 1),
|
||||
)
|
||||
success_msg += "Review the changes and make sure they are as expected (correct indentation, no duplicate lines, etc). Edit the file again if necessary."
|
||||
return ToolExecResult(
|
||||
output=success_msg,
|
||||
)
|
||||
|
||||
# Note: undo_edit method is not implemented in this version as it was removed
|
||||
|
||||
def read_file(self, path: Path):
|
||||
"""Read the content of a file from a given path; raise a ToolError if an error occurs."""
|
||||
try:
|
||||
return path.read_text()
|
||||
except Exception as e:
|
||||
raise ToolError(f"Ran into {e} while trying to read {path}") from None
|
||||
|
||||
def write_file(self, path: Path, file: str):
|
||||
"""Write the content of a file to a given path; raise a ToolError if an error occurs."""
|
||||
try:
|
||||
_ = path.write_text(file)
|
||||
except Exception as e:
|
||||
raise ToolError(f"Ran into {e} while trying to write to {path}") from None
|
||||
|
||||
def _make_output(
|
||||
self,
|
||||
file_content: str,
|
||||
file_descriptor: str,
|
||||
init_line: int = 1,
|
||||
expand_tabs: bool = True,
|
||||
):
|
||||
"""Generate output for the CLI based on the content of a file."""
|
||||
file_content = maybe_truncate(file_content)
|
||||
if expand_tabs:
|
||||
file_content = file_content.expandtabs()
|
||||
file_content = "\n".join(
|
||||
[f"{i + init_line:6}\t{line}" for i, line in enumerate(file_content.split("\n"))]
|
||||
)
|
||||
return (
|
||||
f"Here's the result of running `cat -n` on {file_descriptor}:\n" + file_content + "\n"
|
||||
)
|
||||
|
||||
async def _view_handler(self, arguments: ToolCallArguments, _path: Path) -> ToolExecResult:
|
||||
view_range = arguments.get("view_range", None)
|
||||
if view_range is None:
|
||||
return await self._view(_path, None)
|
||||
if not (isinstance(view_range, list) and all(isinstance(i, int) for i in view_range)):
|
||||
return ToolExecResult(
|
||||
error="Parameter `view_range` should be a list of integers.",
|
||||
error_code=-1,
|
||||
)
|
||||
view_range_int: list[int] = [i for i in view_range if isinstance(i, int)]
|
||||
return await self._view(_path, view_range_int)
|
||||
|
||||
def _create_handler(self, arguments: ToolCallArguments, _path: Path) -> ToolExecResult:
|
||||
file_text = arguments.get("file_text", None)
|
||||
if not isinstance(file_text, str):
|
||||
return ToolExecResult(
|
||||
error="Parameter `file_text` is required and must be a string for command: create",
|
||||
error_code=-1,
|
||||
)
|
||||
self.write_file(_path, file_text)
|
||||
return ToolExecResult(output=f"File created successfully at: {_path}")
|
||||
|
||||
def _str_replace_handler(self, arguments: ToolCallArguments, _path: Path) -> ToolExecResult:
|
||||
old_str = arguments.get("old_str") if "old_str" in arguments else None
|
||||
if not isinstance(old_str, str):
|
||||
return ToolExecResult(
|
||||
error="Parameter `old_str` is required and should be a string for command: str_replace",
|
||||
error_code=-1,
|
||||
)
|
||||
new_str = arguments.get("new_str") if "new_str" in arguments else None
|
||||
if not (new_str is None or isinstance(new_str, str)):
|
||||
return ToolExecResult(
|
||||
error="Parameter `new_str` should be a string or null for command: str_replace",
|
||||
error_code=-1,
|
||||
)
|
||||
return self.str_replace(_path, old_str, new_str)
|
||||
|
||||
def _insert_handler(self, arguments: ToolCallArguments, _path: Path) -> ToolExecResult:
|
||||
insert_line = arguments.get("insert_line") if "insert_line" in arguments else None
|
||||
if not isinstance(insert_line, int):
|
||||
return ToolExecResult(
|
||||
error="Parameter `insert_line` is required and should be integer for command: insert",
|
||||
error_code=-1,
|
||||
)
|
||||
new_str_to_insert = arguments.get("new_str") if "new_str" in arguments else None
|
||||
if not isinstance(new_str_to_insert, str):
|
||||
return ToolExecResult(
|
||||
error="Parameter `new_str` is required for command: insert",
|
||||
error_code=-1,
|
||||
)
|
||||
return self._insert(_path, insert_line, new_str_to_insert)
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
A powerful CLI wrapper for the TextEditorTool that supports sub-commands.
|
||||
"""
|
||||
parser = argparse.ArgumentParser(description="CLI for TextEditorTool.")
|
||||
subparsers = parser.add_subparsers(dest="command", required=True, help="Sub-command help")
|
||||
|
||||
parser_view = subparsers.add_parser("view", help="View a file or directory.")
|
||||
parser_view.add_argument(
|
||||
"--path", required=True, help="Absolute path to the file or directory."
|
||||
)
|
||||
parser_view.add_argument(
|
||||
"--view_range", nargs=2, type=int, help="Line range to view, e.g., 11 12"
|
||||
)
|
||||
|
||||
parser_create = subparsers.add_parser("create", help="Create a new file.")
|
||||
parser_create.add_argument("--path", required=True, help="Absolute path for the new file.")
|
||||
parser_create.add_argument("--file_text", required=True, help="Content of the new file.")
|
||||
|
||||
parser_replace = subparsers.add_parser("str_replace", help="Replace a string in a file.")
|
||||
parser_replace.add_argument("--path", required=True, help="Absolute path to the file.")
|
||||
parser_replace.add_argument("--old_str", required=True, help="The string to be replaced.")
|
||||
parser_replace.add_argument(
|
||||
"--new_str", required=False, default="", help="The string to replace with."
|
||||
)
|
||||
|
||||
parser_insert = subparsers.add_parser("insert", help="Insert a string at a specific line.")
|
||||
parser_insert.add_argument("--path", required=True, help="Absolute path to the file.")
|
||||
parser_insert.add_argument(
|
||||
"--insert_line", type=int, required=True, help="Line number to insert after."
|
||||
)
|
||||
parser_insert.add_argument("--new_str", required=True, help="The string to insert.")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
tool = TextEditorTool()
|
||||
|
||||
arguments = vars(args)
|
||||
|
||||
try:
|
||||
_path = Path(arguments["path"])
|
||||
|
||||
tool.validate_path(args.command, _path)
|
||||
|
||||
if args.command == "view":
|
||||
result = asyncio.run(tool._view_handler(arguments, _path))
|
||||
elif args.command == "create":
|
||||
result = tool._create_handler(arguments, _path)
|
||||
elif args.command == "str_replace":
|
||||
result = tool._str_replace_handler(arguments, _path)
|
||||
elif args.command == "insert":
|
||||
result = tool._insert_handler(arguments, _path)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Sub-command '{args.command}' is not implemented in CLI wrapper."
|
||||
)
|
||||
|
||||
if result.error:
|
||||
print(f"Error: {result.error}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
else:
|
||||
print(result.output)
|
||||
sys.exit(0)
|
||||
|
||||
except Exception as e:
|
||||
print(f"An unexpected error occurred: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,329 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""JSON editing tool for structured JSON file modifications."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import override
|
||||
|
||||
from jsonpath_ng import Fields, Index
|
||||
from jsonpath_ng import parse as jsonpath_parse
|
||||
from jsonpath_ng.exceptions import JSONPathError
|
||||
|
||||
from trae_agent.tools.base import Tool, ToolCallArguments, ToolError, ToolExecResult, ToolParameter
|
||||
|
||||
|
||||
class JSONEditTool(Tool):
|
||||
"""Tool for editing JSON files using JSONPath expressions."""
|
||||
|
||||
def __init__(self, model_provider: str | None = None) -> None:
|
||||
super().__init__(model_provider)
|
||||
|
||||
@override
|
||||
def get_model_provider(self) -> str | None:
|
||||
return self._model_provider
|
||||
|
||||
@override
|
||||
def get_name(self) -> str:
|
||||
return "json_edit_tool"
|
||||
|
||||
@override
|
||||
def get_description(self) -> str:
|
||||
return """Tool for editing JSON files with JSONPath expressions
|
||||
* Supports targeted modifications to JSON structures using JSONPath syntax
|
||||
* Operations: view, set, add, remove
|
||||
* JSONPath examples: '$.users[0].name', '$.config.database.host', '$.items[*].price'
|
||||
* Safe JSON parsing and validation with detailed error messages
|
||||
* Preserves JSON formatting where possible
|
||||
|
||||
Operation details:
|
||||
- `view`: Display JSON content or specific paths
|
||||
- `set`: Update existing values at specified paths
|
||||
- `add`: Add new key-value pairs (for objects) or append to arrays
|
||||
- `remove`: Delete elements at specified paths
|
||||
|
||||
JSONPath syntax supported:
|
||||
- `$` - root element
|
||||
- `.key` - object property access
|
||||
- `[index]` - array index access
|
||||
- `[*]` - all elements in array/object
|
||||
- `..key` - recursive descent (find key at any level)
|
||||
- `[start:end]` - array slicing
|
||||
"""
|
||||
|
||||
@override
|
||||
def get_parameters(self) -> list[ToolParameter]:
|
||||
"""Get the parameters for the JSON edit tool."""
|
||||
return [
|
||||
ToolParameter(
|
||||
name="operation",
|
||||
type="string",
|
||||
description="The operation to perform on the JSON file.",
|
||||
required=True,
|
||||
enum=["view", "set", "add", "remove"],
|
||||
),
|
||||
ToolParameter(
|
||||
name="file_path",
|
||||
type="string",
|
||||
description="The full, ABSOLUTE path to the JSON file to edit. You MUST combine the [Project root path] with the file's relative path to construct this. Relative paths are NOT allowed.",
|
||||
required=True,
|
||||
),
|
||||
ToolParameter(
|
||||
name="json_path",
|
||||
type="string",
|
||||
description="JSONPath expression to specify the target location (e.g., '$.users[0].name', '$.config.database'). Required for set, add, and remove operations. Optional for view to show specific paths.",
|
||||
required=False,
|
||||
),
|
||||
ToolParameter(
|
||||
name="value",
|
||||
type="object",
|
||||
description="The value to set or add. Must be JSON-serializable. Required for set and add operations.",
|
||||
required=False,
|
||||
),
|
||||
ToolParameter(
|
||||
name="pretty_print",
|
||||
type="boolean",
|
||||
description="Whether to format the JSON output with proper indentation. Defaults to true.",
|
||||
required=False,
|
||||
),
|
||||
]
|
||||
|
||||
@override
|
||||
async def execute(self, arguments: ToolCallArguments) -> ToolExecResult:
|
||||
"""Execute the JSON edit operation."""
|
||||
try:
|
||||
operation = str(arguments.get("operation", "")).lower()
|
||||
if not operation:
|
||||
return ToolExecResult(error="Operation parameter is required", error_code=-1)
|
||||
|
||||
file_path_str = str(arguments.get("file_path", ""))
|
||||
if not file_path_str:
|
||||
return ToolExecResult(error="file_path parameter is required", error_code=-1)
|
||||
|
||||
file_path = Path(file_path_str)
|
||||
if not file_path.is_absolute():
|
||||
return ToolExecResult(
|
||||
error=f"File path must be absolute: {file_path}", error_code=-1
|
||||
)
|
||||
|
||||
json_path_arg = arguments.get("json_path")
|
||||
if json_path_arg is not None and not isinstance(json_path_arg, str):
|
||||
return ToolExecResult(error="json_path parameter must be a string.", error_code=-1)
|
||||
|
||||
value = arguments.get("value")
|
||||
|
||||
pretty_print_arg = arguments.get("pretty_print", True)
|
||||
if not isinstance(pretty_print_arg, bool):
|
||||
return ToolExecResult(
|
||||
error="pretty_print parameter must be a boolean.", error_code=-1
|
||||
)
|
||||
|
||||
if operation == "view":
|
||||
return await self._view_json(file_path, json_path_arg, pretty_print_arg)
|
||||
|
||||
if not isinstance(json_path_arg, str):
|
||||
return ToolExecResult(
|
||||
error=f"json_path parameter is required and must be a string for the '{operation}' operation.",
|
||||
error_code=-1,
|
||||
)
|
||||
|
||||
if operation in ["set", "add"]:
|
||||
if value is None:
|
||||
return ToolExecResult(
|
||||
error=f"A 'value' parameter is required for the '{operation}' operation.",
|
||||
error_code=-1,
|
||||
)
|
||||
if operation == "set":
|
||||
return await self._set_json_value(
|
||||
file_path, json_path_arg, value, pretty_print_arg
|
||||
)
|
||||
else: # operation == "add"
|
||||
return await self._add_json_value(
|
||||
file_path, json_path_arg, value, pretty_print_arg
|
||||
)
|
||||
|
||||
if operation == "remove":
|
||||
return await self._remove_json_value(file_path, json_path_arg, pretty_print_arg)
|
||||
|
||||
return ToolExecResult(
|
||||
error=f"Unknown operation: {operation}. Supported operations: view, set, add, remove",
|
||||
error_code=-1,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
return ToolExecResult(error=f"JSON edit tool error: {str(e)}", error_code=-1)
|
||||
|
||||
async def _load_json_file(self, file_path: Path) -> dict | list:
|
||||
"""Load and parse JSON file."""
|
||||
if not file_path.exists():
|
||||
raise ToolError(f"File does not exist: {file_path}")
|
||||
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
content = f.read().strip()
|
||||
if not content:
|
||||
raise ToolError(f"File is empty: {file_path}")
|
||||
return json.loads(content)
|
||||
except json.JSONDecodeError as e:
|
||||
raise ToolError(f"Invalid JSON in file {file_path}: {str(e)}") from e
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error reading file {file_path}: {str(e)}") from e
|
||||
|
||||
async def _save_json_file(
|
||||
self, file_path: Path, data: dict | list, pretty_print: bool = True
|
||||
) -> None:
|
||||
"""Save JSON data to file."""
|
||||
try:
|
||||
with open(file_path, "w", encoding="utf-8") as f:
|
||||
if pretty_print:
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
else:
|
||||
json.dump(data, f, ensure_ascii=False)
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error writing to file {file_path}: {str(e)}") from e
|
||||
|
||||
def _parse_jsonpath(self, json_path_str: str):
|
||||
"""Parse JSONPath expression with error handling."""
|
||||
try:
|
||||
return jsonpath_parse(json_path_str)
|
||||
except JSONPathError as e:
|
||||
raise ToolError(f"Invalid JSONPath expression '{json_path_str}': {str(e)}") from e
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error parsing JSONPath '{json_path_str}': {str(e)}") from e
|
||||
|
||||
async def _view_json(
|
||||
self, file_path: Path, json_path_str: str | None, pretty_print: bool
|
||||
) -> ToolExecResult:
|
||||
"""View JSON file content or specific paths."""
|
||||
data = await self._load_json_file(file_path)
|
||||
|
||||
if json_path_str:
|
||||
jsonpath_expr = self._parse_jsonpath(json_path_str)
|
||||
matches = jsonpath_expr.find(data)
|
||||
|
||||
if not matches:
|
||||
return ToolExecResult(output=f"No matches found for JSONPath: {json_path_str}")
|
||||
|
||||
result_data = [match.value for match in matches]
|
||||
if len(result_data) == 1:
|
||||
result_data = result_data[0]
|
||||
|
||||
if pretty_print:
|
||||
output = json.dumps(result_data, indent=2, ensure_ascii=False)
|
||||
else:
|
||||
output = json.dumps(result_data, ensure_ascii=False)
|
||||
|
||||
return ToolExecResult(output=f"JSONPath '{json_path_str}' matches:\n{output}")
|
||||
else:
|
||||
if pretty_print:
|
||||
output = json.dumps(data, indent=2, ensure_ascii=False)
|
||||
else:
|
||||
output = json.dumps(data, ensure_ascii=False)
|
||||
|
||||
return ToolExecResult(output=f"JSON content of {file_path}:\n{output}")
|
||||
|
||||
async def _set_json_value(
|
||||
self, file_path: Path, json_path_str: str, value, pretty_print: bool
|
||||
) -> ToolExecResult:
|
||||
"""Set value at specified JSONPath."""
|
||||
data = await self._load_json_file(file_path)
|
||||
jsonpath_expr = self._parse_jsonpath(json_path_str)
|
||||
|
||||
matches = jsonpath_expr.find(data)
|
||||
if not matches:
|
||||
return ToolExecResult(
|
||||
error=f"No matches found for JSONPath: {json_path_str}", error_code=-1
|
||||
)
|
||||
|
||||
updated_data = jsonpath_expr.update(data, value)
|
||||
await self._save_json_file(file_path, updated_data, pretty_print)
|
||||
|
||||
match_count = len(matches)
|
||||
return ToolExecResult(
|
||||
output=f"Successfully updated {match_count} location(s) at JSONPath '{json_path_str}' with value: {json.dumps(value)}"
|
||||
)
|
||||
|
||||
async def _add_json_value(
|
||||
self, file_path: Path, json_path_str: str, value, pretty_print: bool
|
||||
) -> ToolExecResult:
|
||||
"""Add value at specified JSONPath."""
|
||||
data = await self._load_json_file(file_path)
|
||||
jsonpath_expr = self._parse_jsonpath(json_path_str)
|
||||
|
||||
parent_path = jsonpath_expr.left
|
||||
target = jsonpath_expr.right
|
||||
|
||||
parent_matches = parent_path.find(data)
|
||||
if not parent_matches:
|
||||
return ToolExecResult(error=f"Parent path not found: {parent_path}", error_code=-1)
|
||||
|
||||
for match in parent_matches:
|
||||
parent_obj = match.value
|
||||
if isinstance(target, Fields):
|
||||
if not isinstance(parent_obj, dict):
|
||||
return ToolExecResult(
|
||||
error=f"Cannot add key to non-object at path: {parent_path}",
|
||||
error_code=-1,
|
||||
)
|
||||
key_to_add = target.fields[0]
|
||||
parent_obj[key_to_add] = value
|
||||
elif isinstance(target, Index):
|
||||
if not isinstance(parent_obj, list):
|
||||
return ToolExecResult(
|
||||
error=f"Cannot add element to non-array at path: {parent_path}",
|
||||
error_code=-1,
|
||||
)
|
||||
index_to_add = target.index
|
||||
parent_obj.insert(index_to_add, value)
|
||||
else:
|
||||
return ToolExecResult(
|
||||
error=f"Unsupported add operation for path type: {type(target)}. Path must end in a key or array index.",
|
||||
error_code=-1,
|
||||
)
|
||||
|
||||
await self._save_json_file(file_path, data, pretty_print)
|
||||
return ToolExecResult(output=f"Successfully added value at JSONPath '{json_path_str}'")
|
||||
|
||||
async def _remove_json_value(
|
||||
self, file_path: Path, json_path_str: str, pretty_print: bool
|
||||
) -> ToolExecResult:
|
||||
"""Remove value at specified JSONPath."""
|
||||
data = await self._load_json_file(file_path)
|
||||
jsonpath_expr = self._parse_jsonpath(json_path_str)
|
||||
|
||||
matches = jsonpath_expr.find(data)
|
||||
if not matches:
|
||||
return ToolExecResult(
|
||||
error=f"No matches found for JSONPath: {json_path_str}", error_code=-1
|
||||
)
|
||||
match_count = len(matches)
|
||||
|
||||
for match in reversed(matches):
|
||||
parent_path = match.full_path.left
|
||||
target = match.path
|
||||
|
||||
parent_matches = parent_path.find(data)
|
||||
if not parent_matches:
|
||||
continue
|
||||
|
||||
for parent_match in parent_matches:
|
||||
parent_obj = parent_match.value
|
||||
try:
|
||||
if isinstance(target, Fields):
|
||||
key_to_remove = target.fields[0]
|
||||
if isinstance(parent_obj, dict) and key_to_remove in parent_obj:
|
||||
del parent_obj[key_to_remove]
|
||||
elif isinstance(target, Index):
|
||||
index_to_remove = target.index
|
||||
if isinstance(parent_obj, list) and -len(
|
||||
parent_obj
|
||||
) <= index_to_remove < len(parent_obj):
|
||||
parent_obj.pop(index_to_remove)
|
||||
except (KeyError, IndexError):
|
||||
pass
|
||||
|
||||
await self._save_json_file(file_path, data, pretty_print)
|
||||
return ToolExecResult(
|
||||
output=f"Successfully removed {match_count} element(s) at JSONPath '{json_path_str}'"
|
||||
)
|
||||
@@ -0,0 +1,282 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from jsonpath_ng import Fields, Index
|
||||
from jsonpath_ng import parse as jsonpath_parse
|
||||
from jsonpath_ng.exceptions import JSONPathError
|
||||
|
||||
|
||||
def override(f):
|
||||
"""A no-op decorator to satisfy the @override syntax."""
|
||||
return f
|
||||
|
||||
|
||||
class Tool:
|
||||
"""A minimal base class to satisfy 'class JSONEditTool(Tool):'."""
|
||||
|
||||
def __init__(self, model_provider: str | None = None) -> None:
|
||||
self._model_provider = model_provider
|
||||
|
||||
|
||||
ToolCallArguments = dict
|
||||
|
||||
|
||||
class ToolError(Exception):
|
||||
"""Custom exception for tool-related errors."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ToolExecResult:
|
||||
"""A class to encapsulate the result of a tool execution."""
|
||||
|
||||
def __init__(self, output: str | None = None, error: str | None = None, error_code: int = 0):
|
||||
self.output = output
|
||||
self.error = error
|
||||
self.error_code = error_code
|
||||
|
||||
|
||||
class ToolParameter:
|
||||
"""A dummy class to allow the get_parameters method to exist without error."""
|
||||
|
||||
def __init__(self, name: str, type: str, description: str, required: bool = False, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
class JSONEditTool(Tool):
|
||||
"""Tool for editing JSON files using JSONPath expressions."""
|
||||
|
||||
def __init__(self, model_provider: str | None = None) -> None:
|
||||
super().__init__(model_provider)
|
||||
|
||||
@override
|
||||
def get_model_provider(self) -> str | None:
|
||||
return self._model_provider
|
||||
|
||||
@override
|
||||
def get_name(self) -> str:
|
||||
return "json_edit_tool"
|
||||
|
||||
@override
|
||||
def get_description(self) -> str:
|
||||
return """..."""
|
||||
|
||||
@override
|
||||
def get_parameters(self) -> list[ToolParameter]:
|
||||
return []
|
||||
|
||||
@override
|
||||
async def execute(self, arguments: ToolCallArguments) -> ToolExecResult:
|
||||
raise NotImplementedError("This method is not used in CLI mode.")
|
||||
|
||||
async def _load_json_file(self, file_path: Path) -> dict | list:
|
||||
if not file_path.exists():
|
||||
raise ToolError(f"File does not exist: {file_path}")
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
content = f.read().strip()
|
||||
if not content:
|
||||
raise ToolError(f"File is empty: {file_path}")
|
||||
return json.loads(content)
|
||||
except json.JSONDecodeError as e:
|
||||
raise ToolError(f"Invalid JSON in file {file_path}: {str(e)}") from e
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error reading file {file_path}: {str(e)}") from e
|
||||
|
||||
async def _save_json_file(
|
||||
self, file_path: Path, data: dict | list, pretty_print: bool = True
|
||||
) -> None:
|
||||
try:
|
||||
with open(file_path, "w", encoding="utf-8") as f:
|
||||
if pretty_print:
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
else:
|
||||
json.dump(data, f, ensure_ascii=False)
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error writing to file {file_path}: {str(e)}") from e
|
||||
|
||||
def _parse_jsonpath(self, json_path_str: str):
|
||||
try:
|
||||
return jsonpath_parse(json_path_str)
|
||||
except JSONPathError as e:
|
||||
raise ToolError(f"Invalid JSONPath expression '{json_path_str}': {str(e)}") from e
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error parsing JSONPath '{json_path_str}': {str(e)}") from e
|
||||
|
||||
async def _view_json(
|
||||
self, file_path: Path, json_path_str: str | None, pretty_print: bool
|
||||
) -> ToolExecResult:
|
||||
data = await self._load_json_file(file_path)
|
||||
if json_path_str:
|
||||
jsonpath_expr = self._parse_jsonpath(json_path_str)
|
||||
matches = jsonpath_expr.find(data)
|
||||
if not matches:
|
||||
return ToolExecResult(output=f"No matches found for JSONPath: {json_path_str}")
|
||||
result_data = [match.value for match in matches]
|
||||
if len(result_data) == 1:
|
||||
result_data = result_data[0]
|
||||
output = json.dumps(result_data, indent=2 if pretty_print else None, ensure_ascii=False)
|
||||
return ToolExecResult(output=f"JSONPath '{json_path_str}' matches:\n{output}")
|
||||
else:
|
||||
output = json.dumps(data, indent=2 if pretty_print else None, ensure_ascii=False)
|
||||
return ToolExecResult(output=f"JSON content of {file_path}:\n{output}")
|
||||
|
||||
async def _set_json_value(
|
||||
self, file_path: Path, json_path_str: str, value, pretty_print: bool
|
||||
) -> ToolExecResult:
|
||||
data = await self._load_json_file(file_path)
|
||||
jsonpath_expr = self._parse_jsonpath(json_path_str)
|
||||
matches = jsonpath_expr.find(data)
|
||||
if not matches:
|
||||
return ToolExecResult(
|
||||
error=f"No matches found for JSONPath: {json_path_str}", error_code=-1
|
||||
)
|
||||
updated_data = jsonpath_expr.update(data, value)
|
||||
await self._save_json_file(file_path, updated_data, pretty_print)
|
||||
return ToolExecResult(
|
||||
output=f"Successfully updated {len(matches)} location(s) at JSONPath '{json_path_str}'"
|
||||
)
|
||||
|
||||
async def _add_json_value(
|
||||
self, file_path: Path, json_path_str: str, value, pretty_print: bool
|
||||
) -> ToolExecResult:
|
||||
data = await self._load_json_file(file_path)
|
||||
jsonpath_expr = self._parse_jsonpath(json_path_str)
|
||||
parent_path, target = jsonpath_expr.left, jsonpath_expr.right
|
||||
parent_matches = parent_path.find(data)
|
||||
if not parent_matches:
|
||||
return ToolExecResult(error=f"Parent path not found: {parent_path}", error_code=-1)
|
||||
for match in parent_matches:
|
||||
parent_obj = match.value
|
||||
if isinstance(target, Fields):
|
||||
if not isinstance(parent_obj, dict):
|
||||
return ToolExecResult(
|
||||
error=f"Cannot add key to non-object at path: {parent_path}", error_code=-1
|
||||
)
|
||||
parent_obj[target.fields[0]] = value
|
||||
elif isinstance(target, Index):
|
||||
if not isinstance(parent_obj, list):
|
||||
return ToolExecResult(
|
||||
error=f"Cannot add element to non-array at path: {parent_path}",
|
||||
error_code=-1,
|
||||
)
|
||||
parent_obj.insert(target.index, value)
|
||||
else:
|
||||
return ToolExecResult(
|
||||
error=f"Unsupported add operation for path type: {type(target)}", error_code=-1
|
||||
)
|
||||
await self._save_json_file(file_path, data, pretty_print)
|
||||
return ToolExecResult(output=f"Successfully added value at JSONPath '{json_path_str}'")
|
||||
|
||||
async def _remove_json_value(
|
||||
self, file_path: Path, json_path_str: str, pretty_print: bool
|
||||
) -> ToolExecResult:
|
||||
data = await self._load_json_file(file_path)
|
||||
jsonpath_expr = self._parse_jsonpath(json_path_str)
|
||||
matches = jsonpath_expr.find(data)
|
||||
if not matches:
|
||||
return ToolExecResult(
|
||||
error=f"No matches found for JSONPath: {json_path_str}", error_code=-1
|
||||
)
|
||||
match_count = len(matches)
|
||||
jsonpath_expr.filter(
|
||||
lambda v: True, data
|
||||
) # This is a conceptual way to remove, actual removal is more complex
|
||||
# A more robust remove logic:
|
||||
for match in reversed(matches):
|
||||
parent_path = match.full_path.left
|
||||
target = match.path
|
||||
for parent_match in parent_path.find(data):
|
||||
parent_obj = parent_match.value
|
||||
try:
|
||||
if isinstance(target, Fields):
|
||||
del parent_obj[target.fields[0]]
|
||||
elif isinstance(target, Index):
|
||||
parent_obj.pop(target.index)
|
||||
except (KeyError, IndexError):
|
||||
pass
|
||||
await self._save_json_file(file_path, data, pretty_print)
|
||||
return ToolExecResult(
|
||||
output=f"Successfully removed {match_count} element(s) at JSONPath '{json_path_str}'"
|
||||
)
|
||||
|
||||
|
||||
async def amain():
|
||||
parser = argparse.ArgumentParser(description="A CLI wrapper for the JSONEditTool.")
|
||||
parser.add_argument(
|
||||
"--operation",
|
||||
required=True,
|
||||
choices=["view", "set", "add", "remove"],
|
||||
help="The operation to perform.",
|
||||
)
|
||||
parser.add_argument("--file_path", required=True, help="Absolute path to the JSON file.")
|
||||
parser.add_argument("--json_path", help="JSONPath expression for the target.")
|
||||
parser.add_argument(
|
||||
"--value",
|
||||
help="The value to set or add, as a JSON string (e.g., '\"a string\"', '123', '{\"key\":\"val\"}').",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pretty_print",
|
||||
type=lambda v: v.lower() == "true",
|
||||
default=True,
|
||||
help="Pretty print the output JSON. Defaults to True.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
tool = JSONEditTool()
|
||||
|
||||
file_path = Path(args.file_path)
|
||||
|
||||
parsed_value = None
|
||||
if args.value is not None:
|
||||
try:
|
||||
parsed_value = json.loads(args.value)
|
||||
except json.JSONDecodeError:
|
||||
print(
|
||||
f"Error: The provided --value is not a valid JSON string: {args.value}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
if not file_path.is_absolute():
|
||||
raise ToolError(f"File path must be absolute: {file_path}")
|
||||
|
||||
result = None
|
||||
if args.operation == "view":
|
||||
result = await tool._view_json(file_path, args.json_path, args.pretty_print)
|
||||
elif args.operation == "set":
|
||||
if args.json_path is None or parsed_value is None:
|
||||
raise ToolError("--json_path and --value are required for 'set' operation.")
|
||||
result = await tool._set_json_value(
|
||||
file_path, args.json_path, parsed_value, args.pretty_print
|
||||
)
|
||||
elif args.operation == "add":
|
||||
if args.json_path is None or parsed_value is None:
|
||||
raise ToolError("--json_path and --value are required for 'add' operation.")
|
||||
result = await tool._add_json_value(
|
||||
file_path, args.json_path, parsed_value, args.pretty_print
|
||||
)
|
||||
elif args.operation == "remove":
|
||||
if args.json_path is None:
|
||||
raise ToolError("--json_path is required for 'remove' operation.")
|
||||
result = await tool._remove_json_value(file_path, args.json_path, args.pretty_print)
|
||||
|
||||
if result.error:
|
||||
print(f"Error: {result.error}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
else:
|
||||
print(result.output)
|
||||
sys.exit(0)
|
||||
|
||||
except ToolError as e:
|
||||
print(f"An error occurred: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(amain())
|
||||
@@ -0,0 +1,58 @@
|
||||
from typing import override
|
||||
|
||||
import mcp
|
||||
|
||||
from .base import Tool, ToolCallArguments, ToolExecResult, ToolParameter
|
||||
|
||||
|
||||
class MCPTool(Tool):
|
||||
def __init__(self, client, tool: mcp.types.Tool, model_provider: str | None = None):
|
||||
super().__init__(model_provider)
|
||||
self.client = client
|
||||
self.tool = tool
|
||||
|
||||
@override
|
||||
def get_model_provider(self) -> str | None:
|
||||
return self._model_provider
|
||||
|
||||
@override
|
||||
def get_name(self) -> str:
|
||||
return self.tool.name
|
||||
|
||||
@override
|
||||
def get_description(self) -> str:
|
||||
return self.tool.description
|
||||
|
||||
@override
|
||||
def get_parameters(self) -> list[ToolParameter]:
|
||||
# For OpenAI models, all parameters must be required=True
|
||||
# For other providers, optional parameters can have required=False
|
||||
def properties_to_parameter():
|
||||
parameters = []
|
||||
inputSchema = self.tool.inputSchema
|
||||
required = inputSchema.get("required", [])
|
||||
properties = inputSchema.get("properties", {})
|
||||
for name, prop in properties.items():
|
||||
tool_para = ToolParameter(
|
||||
name=name,
|
||||
type=prop["type"],
|
||||
items=prop.get("items", None),
|
||||
description=prop["description"],
|
||||
required=name in required,
|
||||
)
|
||||
parameters.append(tool_para)
|
||||
return parameters
|
||||
|
||||
return properties_to_parameter()
|
||||
|
||||
@override
|
||||
async def execute(self, arguments: ToolCallArguments) -> ToolExecResult:
|
||||
try:
|
||||
output = await self.client.call_tool(self.get_name(), arguments)
|
||||
if output.isError:
|
||||
return ToolExecResult(output=None, error=output.content[0].text)
|
||||
else:
|
||||
return ToolExecResult(output=output.content[0].text)
|
||||
|
||||
except Exception as e:
|
||||
return ToolExecResult(error=f"Error running mcp tool: {e}", error_code=-1)
|
||||
@@ -0,0 +1,50 @@
|
||||
# Copyright (c) 2023 Anthropic
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: MIT
|
||||
#
|
||||
# This file has been modified by ByteDance Ltd. and/or its affiliates. on 13 June 2025
|
||||
#
|
||||
# Original file was released under MIT License, with the full license text
|
||||
# available at https://github.com/anthropics/anthropic-quickstarts/blob/main/LICENSE
|
||||
#
|
||||
# This modified file is released under the same license.
|
||||
|
||||
"""Utility to run shell commands asynchronously with a timeout."""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
|
||||
TRUNCATED_MESSAGE: str = "<response clipped><NOTE>To save on context only part of this file has been shown to you. You should retry this tool after you have searched inside the file with `grep -n` in order to find the line numbers of what you are looking for.</NOTE>"
|
||||
MAX_RESPONSE_LEN: int = 16000
|
||||
|
||||
|
||||
def maybe_truncate(content: str, truncate_after: int | None = MAX_RESPONSE_LEN):
|
||||
"""Truncate content and append a notice if content exceeds the specified length."""
|
||||
return (
|
||||
content
|
||||
if not truncate_after or len(content) <= truncate_after
|
||||
else content[:truncate_after] + TRUNCATED_MESSAGE
|
||||
)
|
||||
|
||||
|
||||
async def run(
|
||||
cmd: str,
|
||||
timeout: float | None = 120.0, # seconds
|
||||
truncate_after: int | None = MAX_RESPONSE_LEN,
|
||||
):
|
||||
"""Run a shell command asynchronously with a timeout."""
|
||||
process = await asyncio.create_subprocess_shell(
|
||||
cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
|
||||
try:
|
||||
stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout)
|
||||
return (
|
||||
process.returncode or 0,
|
||||
maybe_truncate(stdout.decode(), truncate_after=truncate_after),
|
||||
maybe_truncate(stderr.decode(), truncate_after=truncate_after),
|
||||
)
|
||||
except asyncio.TimeoutError as exc:
|
||||
with contextlib.suppress(ProcessLookupError):
|
||||
process.kill()
|
||||
raise TimeoutError(f"Command '{cmd}' timed out after {timeout} seconds") from exc
|
||||
@@ -0,0 +1,319 @@
|
||||
# Copyright (c) 2023 Anthropic
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: MIT
|
||||
#
|
||||
# This file has been modified by ByteDance Ltd. and/or its affiliates. on 13 June 2025
|
||||
#
|
||||
# Original file was released under MIT License, with the full license text
|
||||
# available at https://github.com/anthropics/anthropic-quickstarts/blob/main/LICENSE
|
||||
#
|
||||
# This modified file is released under the same license.
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from typing import override
|
||||
|
||||
from trae_agent.tools.base import Tool, ToolCallArguments, ToolExecResult, ToolParameter
|
||||
|
||||
|
||||
@dataclass
|
||||
class ThoughtData:
|
||||
thought: str
|
||||
thought_number: int
|
||||
total_thoughts: int
|
||||
next_thought_needed: bool
|
||||
is_revision: bool | None = None
|
||||
revises_thought: int | None = None
|
||||
branch_from_thought: int | None = None
|
||||
branch_id: str | None = None
|
||||
needs_more_thoughts: bool | None = None
|
||||
|
||||
|
||||
class SequentialThinkingTool(Tool):
|
||||
"""A tool for sequential thinking that helps break down complex problems.
|
||||
|
||||
This tool helps analyze problems through a flexible thinking process that can adapt and evolve.
|
||||
Each thought can build on, question, or revise previous insights as understanding deepens.
|
||||
"""
|
||||
|
||||
@override
|
||||
def get_name(self) -> str:
|
||||
return "sequentialthinking"
|
||||
|
||||
@override
|
||||
def get_description(self) -> str:
|
||||
return """A detailed tool for dynamic and reflective problem-solving through thoughts.
|
||||
This tool helps analyze problems through a flexible thinking process that can adapt and evolve.
|
||||
Each thought can build on, question, or revise previous insights as understanding deepens.
|
||||
|
||||
When to use this tool:
|
||||
- Breaking down complex problems into steps
|
||||
- Planning and design with room for revision
|
||||
- Analysis that might need course correction
|
||||
- Problems where the full scope might not be clear initially
|
||||
- Problems that require a multi-step solution
|
||||
- Tasks that need to maintain context over multiple steps
|
||||
- Situations where irrelevant information needs to be filtered out
|
||||
|
||||
Key features:
|
||||
- You can adjust total_thoughts up or down as you progress
|
||||
- You can question or revise previous thoughts
|
||||
- You can add more thoughts even after reaching what seemed like the end
|
||||
- You can express uncertainty and explore alternative approaches
|
||||
- Not every thought needs to build linearly - you can branch or backtrack
|
||||
- Generates a solution hypothesis
|
||||
- Verifies the hypothesis based on the Chain of Thought steps
|
||||
- Repeats the process until satisfied
|
||||
- Provides a correct answer
|
||||
|
||||
Parameters explained:
|
||||
- thought: Your current thinking step, which can include:
|
||||
* Regular analytical steps
|
||||
* Revisions of previous thoughts
|
||||
* Questions about previous decisions
|
||||
* Realizations about needing more analysis
|
||||
* Changes in approach
|
||||
* Hypothesis generation
|
||||
* Hypothesis verification
|
||||
- next_thought_needed: True if you need more thinking, even if at what seemed like the end
|
||||
- thought_number: Current number in sequence (can go beyond initial total if needed)
|
||||
- total_thoughts: Current estimate of thoughts needed (can be adjusted up/down)
|
||||
- is_revision: A boolean indicating if this thought revises previous thinking
|
||||
- revises_thought: If is_revision is true, which thought number is being reconsidered
|
||||
- branch_from_thought: If branching, which thought number is the branching point
|
||||
- branch_id: Identifier for the current branch (if any)
|
||||
- needs_more_thoughts: If reaching end but realizing more thoughts needed
|
||||
|
||||
You should:
|
||||
1. Start with an initial estimate of needed thoughts, but be ready to adjust
|
||||
2. Feel free to question or revise previous thoughts
|
||||
3. Don't hesitate to add more thoughts if needed, even at the "end"
|
||||
4. Express uncertainty when present
|
||||
5. Mark thoughts that revise previous thinking or branch into new paths
|
||||
6. Ignore information that is irrelevant to the current step
|
||||
7. Generate a solution hypothesis when appropriate
|
||||
8. Verify the hypothesis based on the Chain of Thought steps
|
||||
9. Repeat the process until satisfied with the solution
|
||||
10. Provide a single, ideally correct answer as the final output
|
||||
11. Only set next_thought_needed to false when truly done and a satisfactory answer is reached"""
|
||||
|
||||
@override
|
||||
def get_parameters(self) -> list[ToolParameter]:
|
||||
return [
|
||||
ToolParameter(
|
||||
name="thought",
|
||||
type="string",
|
||||
description="Your current thinking step",
|
||||
required=True,
|
||||
),
|
||||
ToolParameter(
|
||||
name="next_thought_needed",
|
||||
type="boolean",
|
||||
description="Whether another thought step is needed",
|
||||
required=True,
|
||||
),
|
||||
ToolParameter(
|
||||
name="thought_number",
|
||||
type="integer",
|
||||
description="Current thought number. Minimum value is 1.",
|
||||
required=True,
|
||||
),
|
||||
ToolParameter(
|
||||
name="total_thoughts",
|
||||
type="integer",
|
||||
description="Estimated total thoughts needed. Minimum value is 1.",
|
||||
required=True,
|
||||
),
|
||||
ToolParameter(
|
||||
name="is_revision",
|
||||
type="boolean",
|
||||
description="Whether this revises previous thinking",
|
||||
),
|
||||
ToolParameter(
|
||||
name="revises_thought",
|
||||
type="integer",
|
||||
description="Which thought is being reconsidered. Minimum value is 1.",
|
||||
),
|
||||
ToolParameter(
|
||||
name="branch_from_thought",
|
||||
type="integer",
|
||||
description="Branching point thought number. Minimum value is 1.",
|
||||
),
|
||||
ToolParameter(
|
||||
name="branch_id",
|
||||
type="string",
|
||||
description="Branch identifier",
|
||||
),
|
||||
ToolParameter(
|
||||
name="needs_more_thoughts",
|
||||
type="boolean",
|
||||
description="If more thoughts are needed",
|
||||
),
|
||||
]
|
||||
|
||||
def __init__(self, model_provider: str | None = None) -> None:
|
||||
super().__init__(model_provider)
|
||||
self.thought_history: list[ThoughtData] = []
|
||||
self.branches: dict[str, list[ThoughtData]] = {}
|
||||
|
||||
@override
|
||||
def get_model_provider(self) -> str | None:
|
||||
return self._model_provider
|
||||
|
||||
def _validate_thought_data(self, arguments: ToolCallArguments) -> ThoughtData:
|
||||
"""Validate the input arguments and return a ThoughtData object."""
|
||||
if "thought" not in arguments or not isinstance(arguments["thought"], str):
|
||||
raise ValueError("Invalid thought: must be a string")
|
||||
|
||||
if "thought_number" not in arguments or not isinstance(arguments["thought_number"], int):
|
||||
raise ValueError("Invalid thought_number: must be a number")
|
||||
|
||||
if "total_thoughts" not in arguments or not isinstance(arguments["total_thoughts"], int):
|
||||
raise ValueError("Invalid total_thoughts: must be a number")
|
||||
|
||||
if "next_thought_needed" not in arguments or not isinstance(
|
||||
arguments["next_thought_needed"], bool
|
||||
):
|
||||
raise ValueError("Invalid next_thought_needed: must be a boolean")
|
||||
|
||||
# Validate minimum values
|
||||
if arguments["thought_number"] < 1:
|
||||
raise ValueError("thought_number must be at least 1")
|
||||
|
||||
if arguments["total_thoughts"] < 1:
|
||||
raise ValueError("total_thoughts must be at least 1")
|
||||
|
||||
# Validate optional revision fields
|
||||
if (
|
||||
"revises_thought" in arguments
|
||||
and arguments["revises_thought"] is not None
|
||||
and arguments["revises_thought"] != 0
|
||||
):
|
||||
if (
|
||||
not isinstance(arguments["revises_thought"], int)
|
||||
or arguments["revises_thought"] < 1
|
||||
):
|
||||
raise ValueError("revises_thought must be a positive integer")
|
||||
else:
|
||||
revises_thought = int(arguments["revises_thought"])
|
||||
else:
|
||||
revises_thought = None
|
||||
|
||||
if (
|
||||
"branch_from_thought" in arguments
|
||||
and arguments["branch_from_thought"] is not None
|
||||
and arguments["branch_from_thought"] != 0
|
||||
):
|
||||
if (
|
||||
not isinstance(arguments["branch_from_thought"], int)
|
||||
or arguments["branch_from_thought"] < 1
|
||||
):
|
||||
raise ValueError("branch_from_thought must be a positive integer")
|
||||
else:
|
||||
branch_from_thought = int(arguments["branch_from_thought"])
|
||||
else:
|
||||
branch_from_thought = None
|
||||
|
||||
# Extract and cast the validated values
|
||||
thought = str(arguments["thought"])
|
||||
thought_number = int(arguments["thought_number"]) # Already validated as int
|
||||
total_thoughts = int(arguments["total_thoughts"]) # Already validated as int
|
||||
next_thought_needed = bool(arguments["next_thought_needed"]) # Already validated as bool
|
||||
|
||||
# Handle optional fields with proper type checking
|
||||
is_revision = None
|
||||
branch_id = None
|
||||
needs_more_thoughts = None
|
||||
|
||||
if "is_revision" in arguments and arguments["is_revision"] is not None:
|
||||
is_revision = bool(arguments["is_revision"])
|
||||
|
||||
if "branch_id" in arguments and arguments["branch_id"] is not None:
|
||||
branch_id = str(arguments["branch_id"])
|
||||
|
||||
if "needs_more_thoughts" in arguments and arguments["needs_more_thoughts"] is not None:
|
||||
needs_more_thoughts = bool(arguments["needs_more_thoughts"])
|
||||
|
||||
return ThoughtData(
|
||||
thought=thought,
|
||||
thought_number=thought_number,
|
||||
total_thoughts=total_thoughts,
|
||||
next_thought_needed=next_thought_needed,
|
||||
is_revision=is_revision,
|
||||
revises_thought=revises_thought,
|
||||
branch_from_thought=branch_from_thought,
|
||||
branch_id=branch_id,
|
||||
needs_more_thoughts=needs_more_thoughts,
|
||||
)
|
||||
|
||||
def _format_thought(self, thought_data: ThoughtData) -> str:
|
||||
"""Format a thought for display with visual styling."""
|
||||
prefix = ""
|
||||
context = ""
|
||||
|
||||
if thought_data.is_revision:
|
||||
prefix = "🔄 Revision"
|
||||
context = f" (revising thought {thought_data.revises_thought})"
|
||||
elif thought_data.branch_from_thought:
|
||||
prefix = "🌿 Branch"
|
||||
context = (
|
||||
f" (from thought {thought_data.branch_from_thought}, ID: {thought_data.branch_id})"
|
||||
)
|
||||
else:
|
||||
prefix = "💭 Thought"
|
||||
context = ""
|
||||
|
||||
header = f"{prefix} {thought_data.thought_number}/{thought_data.total_thoughts}{context}"
|
||||
border_length = max(len(header), len(thought_data.thought)) + 4
|
||||
border = "─" * border_length
|
||||
|
||||
return f"""
|
||||
┌{border}┐
|
||||
│ {header.ljust(border_length - 2)} │
|
||||
├{border}┤
|
||||
│ {thought_data.thought.ljust(border_length - 2)} │
|
||||
└{border}┘"""
|
||||
|
||||
@override
|
||||
async def execute(self, arguments: ToolCallArguments) -> ToolExecResult:
|
||||
"""Execute the sequential thinking tool."""
|
||||
try:
|
||||
# Validate and extract thought data
|
||||
validated_input = self._validate_thought_data(arguments)
|
||||
|
||||
# Adjust total thoughts if current thought number exceeds it
|
||||
if validated_input.thought_number > validated_input.total_thoughts:
|
||||
validated_input.total_thoughts = validated_input.thought_number
|
||||
|
||||
# Add to thought history
|
||||
self.thought_history.append(validated_input)
|
||||
|
||||
# Handle branching
|
||||
if validated_input.branch_from_thought and validated_input.branch_id:
|
||||
if validated_input.branch_id not in self.branches:
|
||||
self.branches[validated_input.branch_id] = []
|
||||
self.branches[validated_input.branch_id].append(validated_input)
|
||||
|
||||
# Format and display the thought
|
||||
# formatted_thought = self._format_thought(validated_input)
|
||||
# print(formatted_thought, flush=True) # Print to stdout for immediate feedback
|
||||
|
||||
# Prepare response
|
||||
response_data = {
|
||||
"thought_number": validated_input.thought_number,
|
||||
"total_thoughts": validated_input.total_thoughts,
|
||||
"next_thought_needed": validated_input.next_thought_needed,
|
||||
"branches": list(self.branches.keys()),
|
||||
"thought_history_length": len(self.thought_history),
|
||||
}
|
||||
|
||||
return ToolExecResult(
|
||||
output=f"Sequential thinking step completed.\n\nStatus:\n{json.dumps(response_data, indent=2)}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_data = {"error": str(e), "status": "failed"}
|
||||
return ToolExecResult(
|
||||
error=f"Sequential thinking failed: {str(e)}\n\nDetails:\n{json.dumps(error_data, indent=2)}",
|
||||
error_code=-1,
|
||||
)
|
||||
@@ -0,0 +1,33 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from typing import override
|
||||
|
||||
from trae_agent.tools.base import Tool, ToolCallArguments, ToolExecResult, ToolParameter
|
||||
|
||||
|
||||
class TaskDoneTool(Tool):
|
||||
"""Tool to mark a task as done."""
|
||||
|
||||
def __init__(self, model_provider: str | None = None) -> None:
|
||||
super().__init__(model_provider)
|
||||
|
||||
@override
|
||||
def get_model_provider(self) -> str | None:
|
||||
return self._model_provider
|
||||
|
||||
@override
|
||||
def get_name(self) -> str:
|
||||
return "task_done"
|
||||
|
||||
@override
|
||||
def get_description(self) -> str:
|
||||
return "Report the completion of the task. Note that you cannot call this tool before any verification is done. You can write reproduce / test script to verify your solution."
|
||||
|
||||
@override
|
||||
def get_parameters(self) -> list[ToolParameter]:
|
||||
return []
|
||||
|
||||
@override
|
||||
async def execute(self, arguments: ToolCallArguments) -> ToolExecResult:
|
||||
return ToolExecResult(output="Task done.")
|
||||
@@ -0,0 +1,18 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""CLI console module for Trae Agent."""
|
||||
|
||||
from .cli_console import CLIConsole, ConsoleMode, ConsoleType
|
||||
from .console_factory import ConsoleFactory
|
||||
from .rich_console import RichCLIConsole
|
||||
from .simple_console import SimpleCLIConsole
|
||||
|
||||
__all__ = [
|
||||
"CLIConsole",
|
||||
"ConsoleMode",
|
||||
"ConsoleType",
|
||||
"SimpleCLIConsole",
|
||||
"RichCLIConsole",
|
||||
"ConsoleFactory",
|
||||
]
|
||||
@@ -0,0 +1,171 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Base CLI Console classes for Trae Agent."""
|
||||
|
||||
import asyncio
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
|
||||
from rich.panel import Panel
|
||||
from rich.table import Table
|
||||
|
||||
from trae_agent.agent.agent_basics import AgentExecution, AgentStep, AgentStepState
|
||||
from trae_agent.utils.config import LakeviewConfig
|
||||
from trae_agent.utils.lake_view import LakeView
|
||||
|
||||
|
||||
class ConsoleMode(Enum):
|
||||
"""Console operation modes."""
|
||||
|
||||
RUN = "run" # Execute single task and exit
|
||||
INTERACTIVE = "interactive" # Take multiple tasks from user input
|
||||
|
||||
|
||||
class ConsoleType(Enum):
|
||||
"""Available console types."""
|
||||
|
||||
SIMPLE = "simple" # Simple text-based console
|
||||
RICH = "rich" # Rich textual-based console with TUI
|
||||
|
||||
|
||||
AGENT_STATE_INFO = {
|
||||
AgentStepState.THINKING: ("blue", "🤔"),
|
||||
AgentStepState.CALLING_TOOL: ("yellow", "🔧"),
|
||||
AgentStepState.REFLECTING: ("magenta", "💭"),
|
||||
AgentStepState.COMPLETED: ("green", "✅"),
|
||||
AgentStepState.ERROR: ("red", "❌"),
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConsoleStep:
|
||||
"""Represents a console step with its display panel and lakeview information."""
|
||||
|
||||
agent_step: AgentStep
|
||||
agent_step_printed: bool = False
|
||||
lake_view_panel_generator: asyncio.Task[Panel | None] | None = None
|
||||
|
||||
|
||||
class CLIConsole(ABC):
|
||||
"""Base class for CLI console implementations."""
|
||||
|
||||
def __init__(
|
||||
self, mode: ConsoleMode = ConsoleMode.RUN, lakeview_config: LakeviewConfig | None = None
|
||||
):
|
||||
"""Initialize the CLI console.
|
||||
|
||||
Args:
|
||||
config: Configuration object containing settings
|
||||
mode: Console operation mode (run or interactive)
|
||||
"""
|
||||
self.mode: ConsoleMode = mode
|
||||
self.set_lakeview(lakeview_config)
|
||||
self.console_step_history: dict[int, ConsoleStep] = {}
|
||||
self.agent_execution: AgentExecution | None = None
|
||||
|
||||
@abstractmethod
|
||||
async def start(self):
|
||||
"""Start the console display. Should be implemented by subclasses."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def update_status(
|
||||
self, agent_step: AgentStep | None = None, agent_execution: AgentExecution | None = None
|
||||
):
|
||||
"""Update the console with agent status.
|
||||
|
||||
Args:
|
||||
agent_step: Current agent step information
|
||||
agent_execution: Complete agent execution information
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def print_task_details(self, details: dict[str, str]):
|
||||
"""Print initial task configuration details."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def print(self, message: str, color: str = "blue", bold: bool = False):
|
||||
"""Print a message to the console."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_task_input(self) -> str | None:
|
||||
"""Get task input from user (for interactive mode).
|
||||
|
||||
Returns:
|
||||
Task string or None if user wants to exit
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_working_dir_input(self) -> str:
|
||||
"""Get working directory input from user (for interactive mode).
|
||||
|
||||
Returns:
|
||||
Working directory path
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def stop(self):
|
||||
"""Stop the console and cleanup resources."""
|
||||
pass
|
||||
|
||||
def set_lakeview(self, lakeview_config: LakeviewConfig | None = None):
|
||||
"""Set the lakeview configuration for the console."""
|
||||
if lakeview_config:
|
||||
self.lake_view: LakeView | None = LakeView(lakeview_config)
|
||||
else:
|
||||
self.lake_view = None
|
||||
|
||||
|
||||
def generate_agent_step_table(agent_step: AgentStep) -> Table:
|
||||
"""Log an agent step to the console."""
|
||||
color, emoji = AGENT_STATE_INFO.get(agent_step.state, ("white", "❓"))
|
||||
|
||||
# Print the step state in a table
|
||||
table = Table(show_header=False, width=120)
|
||||
table.add_column("Step Number", style="cyan", width=15)
|
||||
table.add_column(f"{agent_step.step_number}", style="green", width=105)
|
||||
|
||||
# Add status row
|
||||
table.add_row(
|
||||
"Status",
|
||||
f"[{color}]{emoji} Step {agent_step.step_number}: {agent_step.state.value.title()}[/{color}]",
|
||||
)
|
||||
|
||||
# Add LLM response row
|
||||
if agent_step.llm_response and agent_step.llm_response.content:
|
||||
table.add_row("LLM Response", f"💬 {agent_step.llm_response.content}")
|
||||
|
||||
# Add tool calls row
|
||||
if agent_step.tool_calls:
|
||||
tool_names = [f"[cyan]{call.name}[/cyan]" for call in agent_step.tool_calls]
|
||||
table.add_row("Tools", f"🔧 {', '.join(tool_names)}")
|
||||
|
||||
for tool_call in agent_step.tool_calls:
|
||||
# Build a tool call table with tool name, arguments and result
|
||||
tool_call_table = Table(show_header=False, width=100)
|
||||
tool_call_table.add_column("Arguments", style="green", width=50)
|
||||
tool_call_table.add_column("Result", style="green", width=50)
|
||||
tool_result_str = ""
|
||||
for tool_result in agent_step.tool_results or []:
|
||||
if tool_result.call_id == tool_call.call_id:
|
||||
tool_result_str = tool_result.result or ""
|
||||
break
|
||||
tool_call_table.add_row(f"{tool_call.arguments}", f"{tool_result_str}")
|
||||
table.add_row(tool_call.name, tool_call_table)
|
||||
|
||||
# Add reflection row
|
||||
if agent_step.reflection:
|
||||
table.add_row("Reflection", f"💭 {agent_step.reflection}")
|
||||
|
||||
# Add error row
|
||||
if agent_step.error:
|
||||
table.add_row("Error", f"❌ {agent_step.error}")
|
||||
|
||||
return table
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Console factory for creating different types of CLI consoles."""
|
||||
|
||||
from trae_agent.utils.config import LakeviewConfig
|
||||
|
||||
from .cli_console import CLIConsole, ConsoleMode, ConsoleType
|
||||
from .rich_console import RichCLIConsole
|
||||
from .simple_console import SimpleCLIConsole
|
||||
|
||||
|
||||
class ConsoleFactory:
|
||||
"""Factory class for creating CLI console instances."""
|
||||
|
||||
@staticmethod
|
||||
def create_console(
|
||||
console_type: ConsoleType,
|
||||
mode: ConsoleMode = ConsoleMode.RUN,
|
||||
lakeview_config: LakeviewConfig | None = None,
|
||||
) -> CLIConsole:
|
||||
"""Create a console instance based on type and mode.
|
||||
|
||||
Args:
|
||||
console_type: Type of console to create (SIMPLE or RICH)
|
||||
mode: Console operation mode (RUN or INTERACTIVE)
|
||||
config: Configuration object
|
||||
|
||||
Returns:
|
||||
CLIConsole instance
|
||||
|
||||
Raises:
|
||||
ValueError: If console_type is not supported
|
||||
"""
|
||||
|
||||
if console_type == ConsoleType.SIMPLE:
|
||||
return SimpleCLIConsole(mode=mode, lakeview_config=lakeview_config)
|
||||
elif console_type == ConsoleType.RICH:
|
||||
return RichCLIConsole(mode=mode, lakeview_config=lakeview_config)
|
||||
|
||||
@staticmethod
|
||||
def get_recommended_console_type(mode: ConsoleMode) -> ConsoleType:
|
||||
"""Get the recommended console type for a given mode.
|
||||
|
||||
Args:
|
||||
mode: Console operation mode
|
||||
|
||||
Returns:
|
||||
Recommended console type
|
||||
"""
|
||||
# Rich console is ideal for interactive mode
|
||||
if mode == ConsoleMode.INTERACTIVE:
|
||||
return ConsoleType.RICH
|
||||
# Simple console works well for run mode
|
||||
else:
|
||||
return ConsoleType.SIMPLE
|
||||
@@ -0,0 +1,367 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Rich CLI Console implementation using Textual TUI."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from typing import override
|
||||
|
||||
from rich.panel import Panel
|
||||
from rich.text import Text
|
||||
from textual import on
|
||||
from textual.app import App, ComposeResult
|
||||
from textual.containers import Container
|
||||
from textual.reactive import reactive
|
||||
from textual.suggester import SuggestFromList
|
||||
from textual.widgets import Footer, Header, Input, RichLog, Static
|
||||
|
||||
from trae_agent.agent.agent_basics import AgentExecution, AgentStep, AgentStepState
|
||||
from trae_agent.utils.cli.cli_console import (
|
||||
AGENT_STATE_INFO,
|
||||
CLIConsole,
|
||||
ConsoleMode,
|
||||
ConsoleStep,
|
||||
generate_agent_step_table,
|
||||
)
|
||||
from trae_agent.utils.config import LakeviewConfig
|
||||
|
||||
|
||||
class TokenDisplay(Static):
|
||||
"""Widget to display real-time token usage."""
|
||||
|
||||
total_tokens: reactive[int] = reactive(0)
|
||||
input_tokens: reactive[int] = reactive(0)
|
||||
output_tokens: reactive[int] = reactive(0)
|
||||
|
||||
@override
|
||||
def render(self) -> Text:
|
||||
"""Render the token display."""
|
||||
if self.total_tokens > 0:
|
||||
return Text(
|
||||
f"Tokens: {self.total_tokens:,} total | "
|
||||
+ f"Input: {self.input_tokens:,} | "
|
||||
+ f"Output: {self.output_tokens:,}",
|
||||
style="bold blue",
|
||||
)
|
||||
return Text("Tokens: 0 total", style="dim")
|
||||
|
||||
def update_tokens(self, agent_execution: AgentExecution):
|
||||
"""Update token counts from agent execution."""
|
||||
if agent_execution and agent_execution.total_tokens:
|
||||
self.input_tokens = agent_execution.total_tokens.input_tokens
|
||||
self.output_tokens = agent_execution.total_tokens.output_tokens
|
||||
self.total_tokens = self.input_tokens + self.output_tokens
|
||||
|
||||
|
||||
class RichConsoleApp(App[None]):
|
||||
"""Textual app for the rich console."""
|
||||
|
||||
CSS_PATH = "rich_console.tcss"
|
||||
|
||||
BINDINGS = [
|
||||
("ctrl+c", "quit", "Quit"),
|
||||
("ctrl+q", "quit", "Quit"),
|
||||
]
|
||||
|
||||
def __init__(self, console_impl: "RichCLIConsole"):
|
||||
super().__init__()
|
||||
self.console_impl: "RichCLIConsole" = console_impl
|
||||
self.execution_log: RichLog | None = None
|
||||
self.task_input: Input | None = None
|
||||
self.task_display: Static | None = None
|
||||
self.token_display: TokenDisplay | None = None
|
||||
self.current_task: str | None = None
|
||||
self.is_running_task: bool = False
|
||||
|
||||
self.options: list[str] = ["help", "exit", "status", "clear"]
|
||||
|
||||
@override
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Compose the UI layout."""
|
||||
yield Header(show_clock=True)
|
||||
|
||||
# Top container for agent execution
|
||||
with Container(id="execution_container"):
|
||||
yield RichLog(id="execution_log", wrap=True, markup=True)
|
||||
|
||||
# Bottom container for input/task display
|
||||
with Container(id="input_container"):
|
||||
if self.console_impl.mode == ConsoleMode.INTERACTIVE:
|
||||
yield Input(
|
||||
placeholder="Enter your task...",
|
||||
id="task_input",
|
||||
suggester=SuggestFromList(self.options, case_sensitive=True),
|
||||
)
|
||||
yield Static("", id="task_display", classes="task_display")
|
||||
else:
|
||||
yield Static("", id="task_display", classes="task_display")
|
||||
|
||||
# Footer container for token usage
|
||||
with Container(id="footer_container"):
|
||||
yield TokenDisplay(id="token_display")
|
||||
|
||||
yield Footer()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Called when the app is mounted."""
|
||||
self.title = "Trae Agent CLI"
|
||||
|
||||
self.execution_log = self.query_one("#execution_log", RichLog)
|
||||
self.token_display = self.query_one("#token_display", TokenDisplay)
|
||||
self.task_display = self.query_one("#task_display", Static)
|
||||
|
||||
if self.console_impl.mode == ConsoleMode.INTERACTIVE:
|
||||
self.task_input = self.query_one("#task_input", Input)
|
||||
_ = self.task_input.focus()
|
||||
|
||||
# Show initial task in RUN mode
|
||||
if self.console_impl.mode == ConsoleMode.RUN and self.console_impl.initial_task:
|
||||
self.task_display.update(
|
||||
Panel(self.console_impl.initial_task, title="Task", border_style="blue")
|
||||
)
|
||||
|
||||
@on(Input.Submitted, "#task_input")
|
||||
def handle_task_input(self, event: Input.Submitted) -> None:
|
||||
"""Handle task input submission in interactive mode."""
|
||||
if self.is_running_task:
|
||||
return
|
||||
|
||||
task = event.value.strip()
|
||||
if not task:
|
||||
return
|
||||
|
||||
handlers: dict = {
|
||||
"exit": self._exit_handler,
|
||||
"quit": self._exit_handler,
|
||||
"help": self._help_handler,
|
||||
"clear": self._clear_handler,
|
||||
"status": self._status_handler,
|
||||
}
|
||||
|
||||
handler = handlers.get(task.lower())
|
||||
if handler:
|
||||
handler(event) if task.lower() not in ["exit", "quit"] else handler()
|
||||
return
|
||||
|
||||
# Execute the task
|
||||
self.current_task = task
|
||||
if self.task_display:
|
||||
_ = self.task_display.update(Panel(task, title="Current Task", border_style="green"))
|
||||
event.input.value = ""
|
||||
self.is_running_task = True
|
||||
|
||||
# Start task execution
|
||||
_ = asyncio.create_task(self._execute_task(task))
|
||||
|
||||
async def _execute_task(self, task: str):
|
||||
"""Execute a task using the agent."""
|
||||
try:
|
||||
if not hasattr(self.console_impl, "agent") or not self.console_impl.agent:
|
||||
if self.execution_log:
|
||||
_ = self.execution_log.write("[red]Error: Agent not available[/red]")
|
||||
return
|
||||
|
||||
# Get working directory
|
||||
working_dir = os.getcwd()
|
||||
if self.console_impl.mode == ConsoleMode.INTERACTIVE:
|
||||
# For interactive mode, we might want to ask for working directory
|
||||
# For now, use current directory
|
||||
pass
|
||||
|
||||
task_args = {
|
||||
"project_path": working_dir,
|
||||
"issue": task,
|
||||
"must_patch": "false",
|
||||
}
|
||||
|
||||
if self.execution_log:
|
||||
_ = self.execution_log.write(f"[blue]Executing task: {task}[/blue]")
|
||||
|
||||
# Execute the task
|
||||
_ = await self.console_impl.agent.run(task, task_args)
|
||||
|
||||
if self.execution_log:
|
||||
_ = self.execution_log.write("[green]Task completed successfully![/green]")
|
||||
|
||||
except Exception as e:
|
||||
if self.execution_log:
|
||||
_ = self.execution_log.write(f"[red]Error executing task: {e}[/red]")
|
||||
finally:
|
||||
self.is_running_task = False
|
||||
if self.console_impl.mode == ConsoleMode.RUN:
|
||||
# In run mode, exit after task completion
|
||||
await asyncio.sleep(1) # Brief pause to show completion
|
||||
_ = self.exit()
|
||||
else:
|
||||
# In interactive mode, clear task display and re-enable input
|
||||
if self.task_display:
|
||||
_ = self.task_display.update("")
|
||||
if self.task_input:
|
||||
_ = self.task_input.focus()
|
||||
|
||||
def log_agent_step(self, agent_step: AgentStep):
|
||||
"""Log an agent step to the execution log."""
|
||||
color, _ = AGENT_STATE_INFO.get(agent_step.state, ("white", "❓"))
|
||||
|
||||
# Create step display
|
||||
step_content = generate_agent_step_table(agent_step)
|
||||
|
||||
if self.execution_log:
|
||||
_ = self.execution_log.write(
|
||||
Panel(step_content, title=f"Step {agent_step.step_number}", border_style=color)
|
||||
)
|
||||
|
||||
def _help_handler(self, event: Input.Submitted):
|
||||
if self.execution_log:
|
||||
self.execution_log.write(
|
||||
Panel(
|
||||
"""[bold]Available Commands:[/bold]
|
||||
|
||||
• Type any task description to execute it
|
||||
• 'status' - Show agent status
|
||||
• 'clear' - Clear the execution log
|
||||
• 'exit' or 'quit' - End the session""",
|
||||
title="Help",
|
||||
border_style="yellow",
|
||||
)
|
||||
)
|
||||
event.input.value = ""
|
||||
|
||||
def _clear_handler(self, event: Input.Submitted):
|
||||
if self.execution_log:
|
||||
_ = self.execution_log.clear()
|
||||
event.input.value = ""
|
||||
|
||||
def _status_handler(self, event: Input.Submitted):
|
||||
if hasattr(self.console_impl, "agent") and self.console_impl.agent:
|
||||
agent_info = getattr(self.console_impl.agent, "agent_config", None)
|
||||
if agent_info and self.execution_log:
|
||||
_ = self.execution_log.write(
|
||||
Panel(
|
||||
f"""[bold]Provider:[/bold] {agent_info.model.model_provider.provider}
|
||||
[bold]Model:[/bold] {agent_info.model.model}
|
||||
[bold]Working Directory:[/bold] {os.getcwd()}""",
|
||||
title="Agent Status",
|
||||
border_style="blue",
|
||||
)
|
||||
)
|
||||
else:
|
||||
if self.execution_log:
|
||||
_ = self.execution_log.write("[yellow]Agent not initialized[/yellow]")
|
||||
event.input.value = ""
|
||||
|
||||
def _exit_handler(self):
|
||||
self.exit()
|
||||
|
||||
async def action_quit(self) -> None:
|
||||
"""Quit the application."""
|
||||
self.console_impl.should_exit = True
|
||||
_ = self.exit()
|
||||
|
||||
|
||||
class RichCLIConsole(CLIConsole):
|
||||
"""Rich CLI console using Textual for TUI interface."""
|
||||
|
||||
def __init__(
|
||||
self, mode: ConsoleMode = ConsoleMode.RUN, lakeview_config: LakeviewConfig | None = None
|
||||
):
|
||||
"""Initialize the rich CLI console."""
|
||||
super().__init__(mode, lakeview_config)
|
||||
self.app: RichConsoleApp | None = None
|
||||
self.should_exit: bool = False
|
||||
self.initial_task: str | None = None
|
||||
self._is_running: bool = False
|
||||
|
||||
# Agent context for interactive mode
|
||||
self.agent = None
|
||||
self.trae_agent_config = None
|
||||
self.config_file = None
|
||||
self.trajectory_file = None
|
||||
|
||||
@override
|
||||
async def start(self):
|
||||
"""Start the rich console application."""
|
||||
# Prevent multiple starts of the same app
|
||||
if self._is_running:
|
||||
return
|
||||
|
||||
self._is_running = True
|
||||
|
||||
try:
|
||||
if self.app is None:
|
||||
self.app = RichConsoleApp(self)
|
||||
|
||||
# Run the textual app
|
||||
await self.app.run_async()
|
||||
finally:
|
||||
self._is_running = False
|
||||
|
||||
@override
|
||||
def update_status(
|
||||
self, agent_step: AgentStep | None = None, agent_execution: AgentExecution | None = None
|
||||
):
|
||||
"""Update the console with agent status."""
|
||||
if agent_step and self.app:
|
||||
if agent_step.step_number not in self.console_step_history:
|
||||
# update step history
|
||||
self.console_step_history[agent_step.step_number] = ConsoleStep(agent_step)
|
||||
|
||||
if (
|
||||
agent_step.state in [AgentStepState.COMPLETED, AgentStepState.ERROR]
|
||||
and not self.console_step_history[agent_step.step_number].agent_step_printed
|
||||
):
|
||||
self.app.log_agent_step(agent_step)
|
||||
self.console_step_history[agent_step.step_number].agent_step_printed = True
|
||||
|
||||
if agent_execution:
|
||||
self.agent_execution = agent_execution
|
||||
if self.app and self.app.token_display:
|
||||
self.app.token_display.update_tokens(agent_execution)
|
||||
|
||||
@override
|
||||
def print_task_details(self, details: dict[str, str]):
|
||||
"""Print initial task configuration details."""
|
||||
if self.app and self.app.execution_log:
|
||||
content = "\n".join([f"[bold]{key}:[/bold] {value}" for key, value in details.items()])
|
||||
_ = self.app.execution_log.write(
|
||||
Panel(content, title="Task Details", border_style="blue")
|
||||
)
|
||||
|
||||
@override
|
||||
def print(self, message: str, color: str = "blue", bold: bool = False):
|
||||
"""Print a message to the console."""
|
||||
if self.app and self.app.execution_log:
|
||||
formatted_message = f"[bold]{message}[/bold]" if bold else message
|
||||
formatted_message = f"[{color}]{formatted_message}[/{color}]"
|
||||
_ = self.app.execution_log.write(formatted_message)
|
||||
|
||||
@override
|
||||
def get_task_input(self) -> str | None:
|
||||
"""Get task input from user (for interactive mode)."""
|
||||
# This method is not used in rich console as input is handled by the TUI
|
||||
return None
|
||||
|
||||
@override
|
||||
def get_working_dir_input(self) -> str:
|
||||
"""Get working directory input from user (for interactive mode)."""
|
||||
# For now, return current directory. Could be enhanced with a dialog
|
||||
return os.getcwd()
|
||||
|
||||
@override
|
||||
def stop(self):
|
||||
"""Stop the console and cleanup resources."""
|
||||
self.should_exit = True
|
||||
if self.app:
|
||||
_ = self.app.exit()
|
||||
|
||||
def set_agent_context(self, agent, trae_agent_config, config_file, trajectory_file) -> None:
|
||||
"""Set the agent context for task execution in interactive mode."""
|
||||
self.agent = agent
|
||||
self.trae_agent_config = trae_agent_config
|
||||
self.config_file = config_file
|
||||
self.trajectory_file = trajectory_file
|
||||
|
||||
def set_initial_task(self, task: str):
|
||||
"""Set the initial task for RUN mode."""
|
||||
self.initial_task = task
|
||||
@@ -0,0 +1,36 @@
|
||||
Screen {
|
||||
layout: vertical;
|
||||
}
|
||||
|
||||
#execution_container {
|
||||
height: 1fr;
|
||||
border: solid $primary;
|
||||
}
|
||||
|
||||
#input_container {
|
||||
height: auto;
|
||||
max-height: 5;
|
||||
border: solid $secondary;
|
||||
}
|
||||
|
||||
#footer_container {
|
||||
height: 1;
|
||||
background: $background 50%;
|
||||
}
|
||||
|
||||
RichLog {
|
||||
scrollbar-size: 1 1;
|
||||
scrollbar-size-horizontal: 1;
|
||||
}
|
||||
|
||||
Input {
|
||||
height: 3;
|
||||
}
|
||||
|
||||
.task_display {
|
||||
background: $surface;
|
||||
color: $text;
|
||||
padding: 1;
|
||||
height: auto;
|
||||
max-height: 3;
|
||||
}
|
||||
@@ -0,0 +1,240 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Simple CLI Console implementation."""
|
||||
|
||||
import asyncio
|
||||
from typing import override
|
||||
|
||||
from rich.console import Console
|
||||
from rich.markdown import Markdown
|
||||
from rich.panel import Panel
|
||||
from rich.table import Table
|
||||
|
||||
from trae_agent.agent.agent_basics import AgentExecution, AgentState, AgentStep, AgentStepState
|
||||
from trae_agent.utils.cli.cli_console import (
|
||||
AGENT_STATE_INFO,
|
||||
CLIConsole,
|
||||
ConsoleMode,
|
||||
ConsoleStep,
|
||||
generate_agent_step_table,
|
||||
)
|
||||
from trae_agent.utils.config import LakeviewConfig
|
||||
|
||||
|
||||
class SimpleCLIConsole(CLIConsole):
|
||||
"""Simple text-based CLI console that prints agent execution trace."""
|
||||
|
||||
def __init__(
|
||||
self, mode: ConsoleMode = ConsoleMode.RUN, lakeview_config: LakeviewConfig | None = None
|
||||
):
|
||||
"""Initialize the simple CLI console.
|
||||
|
||||
Args:
|
||||
config: Configuration object containing lakeview and other settings
|
||||
mode: Console operation mode
|
||||
"""
|
||||
super().__init__(mode, lakeview_config)
|
||||
self.console: Console = Console()
|
||||
|
||||
@override
|
||||
def update_status(
|
||||
self, agent_step: AgentStep | None = None, agent_execution: AgentExecution | None = None
|
||||
):
|
||||
"""Update the console status with new agent step or execution info."""
|
||||
if agent_step:
|
||||
if agent_step.step_number not in self.console_step_history:
|
||||
# update step history
|
||||
self.console_step_history[agent_step.step_number] = ConsoleStep(agent_step)
|
||||
|
||||
if (
|
||||
agent_step.state in [AgentStepState.COMPLETED, AgentStepState.ERROR]
|
||||
and not self.console_step_history[agent_step.step_number].agent_step_printed
|
||||
):
|
||||
self._print_step_update(agent_step, agent_execution)
|
||||
self.console_step_history[agent_step.step_number].agent_step_printed = True
|
||||
|
||||
# If lakeview is enabled, generate lakeview panel in the background
|
||||
if (
|
||||
self.lake_view
|
||||
and not self.console_step_history[
|
||||
agent_step.step_number
|
||||
].lake_view_panel_generator
|
||||
):
|
||||
self.console_step_history[
|
||||
agent_step.step_number
|
||||
].lake_view_panel_generator = asyncio.create_task(
|
||||
self._create_lakeview_step_display(agent_step)
|
||||
)
|
||||
|
||||
self.agent_execution = agent_execution
|
||||
|
||||
@override
|
||||
async def start(self):
|
||||
"""Start the console - wait for completion and then print summary."""
|
||||
while self.agent_execution is None or (
|
||||
self.agent_execution.agent_state != AgentState.COMPLETED
|
||||
and self.agent_execution.agent_state != AgentState.ERROR
|
||||
):
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# Print lakeview summary if enabled
|
||||
if self.lake_view and self.agent_execution:
|
||||
await self._print_lakeview_summary()
|
||||
|
||||
# Print execution summary
|
||||
if self.agent_execution:
|
||||
self._print_execution_summary()
|
||||
|
||||
def _print_step_update(
|
||||
self, agent_step: AgentStep, agent_execution: AgentExecution | None = None
|
||||
):
|
||||
"""Print a step update as it progresses."""
|
||||
|
||||
table = generate_agent_step_table(agent_step)
|
||||
|
||||
if agent_step.llm_usage:
|
||||
table.add_row(
|
||||
"Token Usage",
|
||||
f"Input: {agent_step.llm_usage.input_tokens} Output: {agent_step.llm_usage.output_tokens}",
|
||||
)
|
||||
|
||||
if agent_execution and agent_execution.total_tokens:
|
||||
table.add_row(
|
||||
"Total Tokens",
|
||||
f"Input: {agent_execution.total_tokens.input_tokens} Output: {agent_execution.total_tokens.output_tokens}",
|
||||
)
|
||||
|
||||
self.console.print(table)
|
||||
|
||||
async def _print_lakeview_summary(self):
|
||||
"""Print lakeview summary of all completed steps."""
|
||||
self.console.print("\n" + "=" * 60)
|
||||
self.console.print("[bold cyan]Lakeview Summary[/bold cyan]")
|
||||
self.console.print("=" * 60)
|
||||
|
||||
for step in self.console_step_history.values():
|
||||
if step.lake_view_panel_generator:
|
||||
lake_view_panel = await step.lake_view_panel_generator
|
||||
if lake_view_panel:
|
||||
self.console.print(lake_view_panel)
|
||||
|
||||
def _print_execution_summary(self):
|
||||
"""Print the final execution summary."""
|
||||
if not self.agent_execution:
|
||||
return
|
||||
|
||||
self.console.print("\n" + "=" * 60)
|
||||
self.console.print("[bold green]Execution Summary[/bold green]")
|
||||
self.console.print("=" * 60)
|
||||
|
||||
# Create summary table
|
||||
table = Table(show_header=False, width=60)
|
||||
table.add_column("Metric", style="cyan", width=20)
|
||||
table.add_column("Value", style="green", width=40)
|
||||
|
||||
table.add_row(
|
||||
"Task",
|
||||
self.agent_execution.task[:50] + "..."
|
||||
if len(self.agent_execution.task) > 50
|
||||
else self.agent_execution.task,
|
||||
)
|
||||
table.add_row("Success", "✅ Yes" if self.agent_execution.success else "❌ No")
|
||||
table.add_row("Steps", str(len(self.agent_execution.steps)))
|
||||
table.add_row("Execution Time", f"{self.agent_execution.execution_time:.2f}s")
|
||||
|
||||
if self.agent_execution.total_tokens:
|
||||
total_tokens = (
|
||||
self.agent_execution.total_tokens.input_tokens
|
||||
+ self.agent_execution.total_tokens.output_tokens
|
||||
)
|
||||
table.add_row("Total Tokens", str(total_tokens))
|
||||
table.add_row("Input Tokens", str(self.agent_execution.total_tokens.input_tokens))
|
||||
table.add_row("Output Tokens", str(self.agent_execution.total_tokens.output_tokens))
|
||||
|
||||
self.console.print(table)
|
||||
|
||||
# Display final result
|
||||
if self.agent_execution.final_result:
|
||||
self.console.print(
|
||||
Panel(
|
||||
Markdown(self.agent_execution.final_result),
|
||||
title="Final Result",
|
||||
border_style="green" if self.agent_execution.success else "red",
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def print_task_details(self, details: dict[str, str]):
|
||||
"""Print initial task configuration details."""
|
||||
renderable = ""
|
||||
for key, value in details.items():
|
||||
renderable += f"[bold]{key}:[/bold] {value}\n"
|
||||
renderable = renderable.strip()
|
||||
self.console.print(
|
||||
Panel(
|
||||
renderable,
|
||||
title="Task Details",
|
||||
border_style="blue",
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def print(self, message: str, color: str = "blue", bold: bool = False):
|
||||
"""Print a message to the console."""
|
||||
message = f"[bold]{message}[/bold]" if bold else message
|
||||
message = f"[{color}]{message}[/{color}]"
|
||||
self.console.print(message)
|
||||
|
||||
@override
|
||||
def get_task_input(self) -> str | None:
|
||||
"""Get task input from user (for interactive mode)."""
|
||||
if self.mode != ConsoleMode.INTERACTIVE:
|
||||
return None
|
||||
|
||||
self.console.print("\n[bold blue]Task:[/bold blue] ", end="")
|
||||
try:
|
||||
task = input()
|
||||
if task.lower() in ["exit", "quit"]:
|
||||
return None
|
||||
return task
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
return None
|
||||
|
||||
@override
|
||||
def get_working_dir_input(self) -> str:
|
||||
"""Get working directory input from user (for interactive mode)."""
|
||||
if self.mode != ConsoleMode.INTERACTIVE:
|
||||
return ""
|
||||
|
||||
self.console.print("[bold blue]Working Directory:[/bold blue] ", end="")
|
||||
try:
|
||||
return input()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
return ""
|
||||
|
||||
@override
|
||||
def stop(self):
|
||||
"""Stop the console and cleanup resources."""
|
||||
# Simple console doesn't need explicit cleanup
|
||||
pass
|
||||
|
||||
async def _create_lakeview_step_display(self, agent_step: AgentStep) -> Panel | None:
|
||||
"""Create lakeview display for a step."""
|
||||
if self.lake_view is None:
|
||||
return None
|
||||
|
||||
lake_view_step = await self.lake_view.create_lakeview_step(agent_step)
|
||||
|
||||
if lake_view_step is None:
|
||||
return None
|
||||
|
||||
color, _ = AGENT_STATE_INFO.get(agent_step.state, ("white", "❓"))
|
||||
|
||||
return Panel(
|
||||
f"""[{lake_view_step.tags_emoji}] The agent [bold]{lake_view_step.desc_task}[/bold]
|
||||
{lake_view_step.desc_details}""",
|
||||
title=f"Step {agent_step.step_number} (Lakeview)",
|
||||
border_style=color,
|
||||
width=80,
|
||||
)
|
||||
@@ -0,0 +1,410 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import yaml
|
||||
|
||||
from trae_agent.utils.legacy_config import LegacyConfig
|
||||
|
||||
|
||||
class ConfigError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelProvider:
|
||||
"""
|
||||
Model provider configuration. For official model providers such as OpenAI and Anthropic,
|
||||
the base_url is optional. api_version is required for Azure.
|
||||
"""
|
||||
|
||||
api_key: str
|
||||
provider: str
|
||||
base_url: str | None = None
|
||||
api_version: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelConfig:
|
||||
"""
|
||||
Model configuration.
|
||||
"""
|
||||
|
||||
model: str
|
||||
model_provider: ModelProvider
|
||||
temperature: float
|
||||
top_p: float
|
||||
top_k: int
|
||||
parallel_tool_calls: bool
|
||||
max_retries: int
|
||||
max_tokens: int | None = None # Legacy max_tokens parameter, optional
|
||||
supports_tool_calling: bool = True
|
||||
candidate_count: int | None = None # Gemini specific field
|
||||
stop_sequences: list[str] | None = None
|
||||
max_completion_tokens: int | None = None # Azure OpenAI specific field
|
||||
|
||||
def get_max_tokens_param(self) -> int:
|
||||
"""Get the maximum tokens parameter value.Prioritizes max_completion_tokens, falls back to max_tokens if not available."""
|
||||
if self.max_completion_tokens is not None:
|
||||
return self.max_completion_tokens
|
||||
elif self.max_tokens is not None:
|
||||
return self.max_tokens
|
||||
else:
|
||||
# Return default value if neither is set
|
||||
return 4096
|
||||
|
||||
def should_use_max_completion_tokens(self) -> bool:
|
||||
"""Determine whether to use the max_completion_tokens parameter.Primarily used for Azure OpenAI's newer models (e.g., gpt-5)."""
|
||||
return (
|
||||
self.max_completion_tokens is not None
|
||||
and self.model_provider.provider == "azure"
|
||||
and ("gpt-5" in self.model or "o3" in self.model or "o4-mini" in self.model)
|
||||
)
|
||||
|
||||
def resolve_config_values(
|
||||
self,
|
||||
*,
|
||||
model_providers: dict[str, ModelProvider] | None = None,
|
||||
provider: str | None = None,
|
||||
model: str | None = None,
|
||||
model_base_url: str | None = None,
|
||||
api_key: str | None = None,
|
||||
):
|
||||
"""
|
||||
When some config values are provided through CLI or environment variables,
|
||||
they will override the values in the config file.
|
||||
"""
|
||||
self.model = str(resolve_config_value(cli_value=model, config_value=self.model))
|
||||
|
||||
# If the user wants to change the model provider, they should either:
|
||||
# * Make sure the provider name is available in the model_providers dict;
|
||||
# * If not, base url and api key should be provided to register a new model provider.
|
||||
if provider:
|
||||
if model_providers and provider in model_providers:
|
||||
self.model_provider = model_providers[provider]
|
||||
elif api_key is None:
|
||||
raise ConfigError("To register a new model provider, an api_key should be provided")
|
||||
else:
|
||||
self.model_provider = ModelProvider(
|
||||
api_key=api_key,
|
||||
provider=provider,
|
||||
base_url=model_base_url,
|
||||
)
|
||||
|
||||
# Map providers to their environment variable names
|
||||
env_var_api_key = str(self.model_provider.provider).upper() + "_API_KEY"
|
||||
env_var_api_base_url = str(self.model_provider.provider).upper() + "_BASE_URL"
|
||||
|
||||
resolved_api_key = resolve_config_value(
|
||||
cli_value=api_key,
|
||||
config_value=self.model_provider.api_key,
|
||||
env_var=env_var_api_key,
|
||||
)
|
||||
|
||||
resolved_api_base_url = resolve_config_value(
|
||||
cli_value=model_base_url,
|
||||
config_value=self.model_provider.base_url,
|
||||
env_var=env_var_api_base_url,
|
||||
)
|
||||
|
||||
if resolved_api_key:
|
||||
self.model_provider.api_key = str(resolved_api_key)
|
||||
|
||||
if resolved_api_base_url:
|
||||
self.model_provider.base_url = str(resolved_api_base_url)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MCPServerConfig:
|
||||
# For stdio transport
|
||||
command: str | None = None
|
||||
args: list[str] | None = None
|
||||
env: dict[str, str] | None = None
|
||||
cwd: str | None = None
|
||||
|
||||
# For sse transport
|
||||
url: str | None = None
|
||||
|
||||
# For streamable http transport
|
||||
http_url: str | None = None
|
||||
headers: dict[str, str] | None = None
|
||||
|
||||
# For websocket transport
|
||||
tcp: str | None = None
|
||||
|
||||
# Common
|
||||
timeout: int | None = None
|
||||
trust: bool | None = None
|
||||
|
||||
# Metadata
|
||||
description: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentConfig:
|
||||
"""
|
||||
Base class for agent configurations.
|
||||
"""
|
||||
|
||||
allow_mcp_servers: list[str]
|
||||
mcp_servers_config: dict[str, MCPServerConfig]
|
||||
max_steps: int
|
||||
model: ModelConfig
|
||||
tools: list[str]
|
||||
|
||||
|
||||
@dataclass
|
||||
class TraeAgentConfig(AgentConfig):
|
||||
"""
|
||||
Trae agent configuration.
|
||||
"""
|
||||
|
||||
enable_lakeview: bool = True
|
||||
tools: list[str] = field(
|
||||
default_factory=lambda: [
|
||||
"bash",
|
||||
"str_replace_based_edit_tool",
|
||||
"sequentialthinking",
|
||||
"task_done",
|
||||
]
|
||||
)
|
||||
|
||||
def resolve_config_values(
|
||||
self,
|
||||
*,
|
||||
max_steps: int | None = None,
|
||||
):
|
||||
resolved_value = resolve_config_value(cli_value=max_steps, config_value=self.max_steps)
|
||||
if resolved_value:
|
||||
self.max_steps = int(resolved_value)
|
||||
|
||||
|
||||
@dataclass
|
||||
class LakeviewConfig:
|
||||
"""
|
||||
Lakeview configuration.
|
||||
"""
|
||||
|
||||
model: ModelConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
"""
|
||||
Configuration class for agents, models and model providers.
|
||||
"""
|
||||
|
||||
lakeview: LakeviewConfig | None = None
|
||||
model_providers: dict[str, ModelProvider] | None = None
|
||||
models: dict[str, ModelConfig] | None = None
|
||||
|
||||
trae_agent: TraeAgentConfig | None = None
|
||||
|
||||
@classmethod
|
||||
def create(
|
||||
cls,
|
||||
*,
|
||||
config_file: str | None = None,
|
||||
config_string: str | None = None,
|
||||
) -> "Config":
|
||||
if config_file and config_string:
|
||||
raise ConfigError("Only one of config_file or config_string should be provided")
|
||||
|
||||
# Parse YAML config from file or string
|
||||
try:
|
||||
if config_file is not None:
|
||||
if config_file.endswith(".json"):
|
||||
return cls.create_from_legacy_config(config_file=config_file)
|
||||
with open(config_file, "r") as f:
|
||||
yaml_config = yaml.safe_load(f)
|
||||
elif config_string is not None:
|
||||
yaml_config = yaml.safe_load(config_string)
|
||||
else:
|
||||
raise ConfigError("No config file or config string provided")
|
||||
except yaml.YAMLError as e:
|
||||
raise ConfigError(f"Error parsing YAML config: {e}") from e
|
||||
|
||||
config = cls()
|
||||
|
||||
# Parse model providers
|
||||
model_providers = yaml_config.get("model_providers", None)
|
||||
if model_providers is not None and len(model_providers.keys()) > 0:
|
||||
config_model_providers: dict[str, ModelProvider] = {}
|
||||
for model_provider_name, model_provider_config in model_providers.items():
|
||||
config_model_providers[model_provider_name] = ModelProvider(**model_provider_config)
|
||||
config.model_providers = config_model_providers
|
||||
else:
|
||||
raise ConfigError("No model providers provided")
|
||||
|
||||
# Parse models and populate model_provider fields
|
||||
models = yaml_config.get("models", None)
|
||||
if models is not None and len(models.keys()) > 0:
|
||||
config_models: dict[str, ModelConfig] = {}
|
||||
for model_name, model_config in models.items():
|
||||
if model_config["model_provider"] not in config_model_providers:
|
||||
raise ConfigError(f"Model provider {model_config['model_provider']} not found")
|
||||
config_models[model_name] = ModelConfig(**model_config)
|
||||
config_models[model_name].model_provider = config_model_providers[
|
||||
model_config["model_provider"]
|
||||
]
|
||||
config.models = config_models
|
||||
else:
|
||||
raise ConfigError("No models provided")
|
||||
|
||||
# Parse lakeview config
|
||||
lakeview = yaml_config.get("lakeview", None)
|
||||
if lakeview is not None:
|
||||
lakeview_model_name = lakeview.get("model", None)
|
||||
if lakeview_model_name is None:
|
||||
raise ConfigError("No model provided for lakeview")
|
||||
lakeview_model = config_models[lakeview_model_name]
|
||||
config.lakeview = LakeviewConfig(
|
||||
model=lakeview_model,
|
||||
)
|
||||
else:
|
||||
config.lakeview = None
|
||||
|
||||
mcp_servers_config = {
|
||||
k: MCPServerConfig(**v) for k, v in yaml_config.get("mcp_servers", {}).items()
|
||||
}
|
||||
allow_mcp_servers = yaml_config.get("allow_mcp_servers", [])
|
||||
|
||||
# Parse agents
|
||||
agents = yaml_config.get("agents", None)
|
||||
if agents is not None and len(agents.keys()) > 0:
|
||||
for agent_name, agent_config in agents.items():
|
||||
agent_model_name = agent_config.get("model", None)
|
||||
if agent_model_name is None:
|
||||
raise ConfigError(f"No model provided for {agent_name}")
|
||||
try:
|
||||
agent_model = config_models[agent_model_name]
|
||||
except KeyError as e:
|
||||
raise ConfigError(f"Model {agent_model_name} not found") from e
|
||||
match agent_name:
|
||||
case "trae_agent":
|
||||
trae_agent_config = TraeAgentConfig(
|
||||
**agent_config,
|
||||
mcp_servers_config=mcp_servers_config,
|
||||
allow_mcp_servers=allow_mcp_servers,
|
||||
)
|
||||
trae_agent_config.model = agent_model
|
||||
if trae_agent_config.enable_lakeview and config.lakeview is None:
|
||||
raise ConfigError("Lakeview is enabled but no lakeview config provided")
|
||||
config.trae_agent = trae_agent_config
|
||||
case _:
|
||||
raise ConfigError(f"Unknown agent: {agent_name}")
|
||||
else:
|
||||
raise ConfigError("No agent configs provided")
|
||||
return config
|
||||
|
||||
def resolve_config_values(
|
||||
self,
|
||||
*,
|
||||
provider: str | None = None,
|
||||
model: str | None = None,
|
||||
model_base_url: str | None = None,
|
||||
api_key: str | None = None,
|
||||
max_steps: int | None = None,
|
||||
):
|
||||
if self.trae_agent:
|
||||
self.trae_agent.resolve_config_values(
|
||||
max_steps=max_steps,
|
||||
)
|
||||
self.trae_agent.model.resolve_config_values(
|
||||
model_providers=self.model_providers,
|
||||
provider=provider,
|
||||
model=model,
|
||||
model_base_url=model_base_url,
|
||||
api_key=api_key,
|
||||
)
|
||||
return self
|
||||
|
||||
@classmethod
|
||||
def create_from_legacy_config(
|
||||
cls,
|
||||
*,
|
||||
legacy_config: LegacyConfig | None = None,
|
||||
config_file: str | None = None,
|
||||
) -> "Config":
|
||||
if legacy_config and config_file:
|
||||
raise ConfigError("Only one of legacy_config or config_file should be provided")
|
||||
|
||||
if config_file:
|
||||
legacy_config = LegacyConfig(config_file)
|
||||
elif not legacy_config:
|
||||
raise ConfigError("No legacy_config or config_file provided")
|
||||
|
||||
model_provider = ModelProvider(
|
||||
api_key=legacy_config.model_providers[legacy_config.default_provider].api_key,
|
||||
base_url=legacy_config.model_providers[legacy_config.default_provider].base_url,
|
||||
api_version=legacy_config.model_providers[legacy_config.default_provider].api_version,
|
||||
provider=legacy_config.default_provider,
|
||||
)
|
||||
|
||||
model_config = ModelConfig(
|
||||
model=legacy_config.model_providers[legacy_config.default_provider].model,
|
||||
model_provider=model_provider,
|
||||
max_tokens=legacy_config.model_providers[legacy_config.default_provider].max_tokens,
|
||||
temperature=legacy_config.model_providers[legacy_config.default_provider].temperature,
|
||||
top_p=legacy_config.model_providers[legacy_config.default_provider].top_p,
|
||||
top_k=legacy_config.model_providers[legacy_config.default_provider].top_k,
|
||||
parallel_tool_calls=legacy_config.model_providers[
|
||||
legacy_config.default_provider
|
||||
].parallel_tool_calls,
|
||||
max_retries=legacy_config.model_providers[legacy_config.default_provider].max_retries,
|
||||
candidate_count=legacy_config.model_providers[
|
||||
legacy_config.default_provider
|
||||
].candidate_count,
|
||||
stop_sequences=legacy_config.model_providers[
|
||||
legacy_config.default_provider
|
||||
].stop_sequences,
|
||||
)
|
||||
mcp_servers_config = {
|
||||
k: MCPServerConfig(**vars(v)) for k, v in legacy_config.mcp_servers.items()
|
||||
}
|
||||
trae_agent_config = TraeAgentConfig(
|
||||
max_steps=legacy_config.max_steps,
|
||||
enable_lakeview=legacy_config.enable_lakeview,
|
||||
model=model_config,
|
||||
allow_mcp_servers=legacy_config.allow_mcp_servers,
|
||||
mcp_servers_config=mcp_servers_config,
|
||||
)
|
||||
|
||||
if trae_agent_config.enable_lakeview:
|
||||
lakeview_config = LakeviewConfig(
|
||||
model=model_config,
|
||||
)
|
||||
else:
|
||||
lakeview_config = None
|
||||
|
||||
return cls(
|
||||
trae_agent=trae_agent_config,
|
||||
lakeview=lakeview_config,
|
||||
model_providers={
|
||||
legacy_config.default_provider: model_provider,
|
||||
},
|
||||
models={
|
||||
"default_model": model_config,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def resolve_config_value(
|
||||
*,
|
||||
cli_value: int | str | float | None,
|
||||
config_value: int | str | float | None,
|
||||
env_var: str | None = None,
|
||||
) -> int | str | float | None:
|
||||
"""Resolve configuration value with priority: CLI > ENV > Config > Default."""
|
||||
if cli_value is not None:
|
||||
return cli_value
|
||||
|
||||
if env_var and os.getenv(env_var):
|
||||
return os.getenv(env_var)
|
||||
|
||||
if config_value is not None:
|
||||
return config_value
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
LOCAL_STORAGE_PATH = Path.home() / ".trae-agent"
|
||||
@@ -0,0 +1,209 @@
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
from trae_agent.agent.agent_basics import AgentStep
|
||||
from trae_agent.utils.config import LakeviewConfig
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage
|
||||
from trae_agent.utils.llm_clients.llm_client import LLMClient
|
||||
|
||||
StepType = tuple[
|
||||
str, # content for human (will write into result file)
|
||||
str
|
||||
| None, # content for llm, or None if no need to analyze (i.e., minor step), watch out length limit
|
||||
]
|
||||
|
||||
|
||||
EXTRACTOR_PROMPT = """
|
||||
Given the preceding excerpt, your job is to determine "what task is the agent performing in <this_step>".
|
||||
Output your answer in two granularities: <task>...</task><details>...</details>.
|
||||
In the <task> tag, the answer should be concise and general. It should omit ANY bug-specific details, and contain at most 10 words.
|
||||
In the <details> tag, the answer should complement the <task> tag by adding bug-specific details. It should be informative and contain at most 30 words.
|
||||
|
||||
Examples:
|
||||
|
||||
<task>The agent is writing a reproduction test script.</task><details>The agent is writing "test_bug.py" to reproduce the bug in XXX-Project's create_foo method not comparing sizes correctly.</details>
|
||||
<task>The agent is examining source code.</task><details>The agent is searching for "function_name" in the code repository, that is related to the "foo.py:function_name" line in the stack trace.</details>
|
||||
<task>The agent is fixing the reproduction test script.</task><details>The agent is fixing "test_bug.py" that forgets to import the function "foo", causing a NameError.</details>
|
||||
|
||||
Now, answer the question "what task is the agent performing in <this_step>".
|
||||
Again, provide only the answer with no other commentary. The format should be "<task>...</task><details>...</details>".
|
||||
"""
|
||||
|
||||
TAGGER_PROMPT = """
|
||||
Given the trajectory, your job is to determine "what task is the agent performing in the current step".
|
||||
Output your answer by choosing the applicable tags in the below list for the current step.
|
||||
If it is performing multiple tasks in one step, choose ALL applicable tags, separated by a comma.
|
||||
|
||||
<tags>
|
||||
WRITE_TEST: It writes a test script to reproduce the bug, or modifies a non-working test script to fix problems found in testing.
|
||||
VERIFY_TEST: It runs the reproduction test script to verify the testing environment is working.
|
||||
EXAMINE_CODE: It views, searches, or explores the code repository to understand the cause of the bug.
|
||||
WRITE_FIX: It modifies the source code to fix the identified bug.
|
||||
VERIFY_FIX: It runs the reproduction test or existing tests to verify the fix indeed solves the bug.
|
||||
REPORT: It reports to the user that the job is completed or some progress has been made.
|
||||
THINK: It analyzes the bug through thinking, but does not perform concrete actions right now.
|
||||
OUTLIER: A major part in this step does not fit into any tag above, such as running a shell command to install dependencies.
|
||||
</tags>
|
||||
|
||||
<examples>
|
||||
If the agent is opening a file to examine, output <tags>EXAMINE_CODE</tags>.
|
||||
If the agent is fixing a known problem in the reproduction test script and then running it again, output <tags>WRITE_TEST,VERIFY_TEST</tags>.
|
||||
If the agent is merely thinking about the root cause of the bug without other actions, output <tags>THINK</tags>.
|
||||
</examples>
|
||||
|
||||
Output only the tags with no other commentary. The format should be <tags>...</tags>
|
||||
"""
|
||||
|
||||
KNOWN_TAGS = {
|
||||
"WRITE_TEST": "☑️",
|
||||
"VERIFY_TEST": "✅",
|
||||
"EXAMINE_CODE": "👁️",
|
||||
"WRITE_FIX": "📝",
|
||||
"VERIFY_FIX": "🔥",
|
||||
"REPORT": "📣",
|
||||
"THINK": "🧠",
|
||||
"OUTLIER": "⁉️",
|
||||
}
|
||||
|
||||
tags_re = re.compile(r"<tags>([A-Z_,\s]+)</tags>")
|
||||
|
||||
|
||||
@dataclass
|
||||
class LakeViewStep:
|
||||
desc_task: str
|
||||
desc_details: str
|
||||
tags_emoji: str
|
||||
|
||||
|
||||
class LakeView:
|
||||
def __init__(self, lake_view_config: LakeviewConfig | None):
|
||||
if lake_view_config is None:
|
||||
return
|
||||
|
||||
self.model_config = lake_view_config.model
|
||||
self.lakeview_llm_client: LLMClient = LLMClient(self.model_config)
|
||||
|
||||
self.steps: list[str] = []
|
||||
|
||||
def get_label(self, tags: None | list[str], emoji: bool = True) -> str:
|
||||
if not tags:
|
||||
return ""
|
||||
|
||||
return " · ".join([KNOWN_TAGS[tag] + tag if emoji else tag for tag in tags])
|
||||
|
||||
async def extract_task_in_step(self, prev_step: str, this_step: str) -> tuple[str, str]:
|
||||
llm_messages = [
|
||||
LLMMessage(
|
||||
role="user",
|
||||
content=f"The following is an excerpt of the steps trying to solve a software bug by an AI agent: <previous_step>{prev_step}</previous_step><this_step>{this_step}</this_step>",
|
||||
),
|
||||
LLMMessage(role="assistant", content="I understand."),
|
||||
LLMMessage(role="user", content=EXTRACTOR_PROMPT),
|
||||
LLMMessage(
|
||||
role="assistant",
|
||||
content="Sure. Here is the task the agent is performing: <task>The agent",
|
||||
),
|
||||
]
|
||||
|
||||
self.model_config.temperature = 0.1
|
||||
llm_response = self.lakeview_llm_client.chat(
|
||||
model_config=self.model_config,
|
||||
messages=llm_messages,
|
||||
reuse_history=False,
|
||||
)
|
||||
|
||||
content = llm_response.content.strip()
|
||||
|
||||
retry = 0
|
||||
while retry < 10 and (
|
||||
"</task>" not in content or "<details>" not in content or "</details>" not in content
|
||||
):
|
||||
retry += 1
|
||||
llm_response = self.lakeview_llm_client.chat(
|
||||
model_config=self.model_config,
|
||||
messages=llm_messages,
|
||||
reuse_history=False,
|
||||
)
|
||||
content = llm_response.content.strip()
|
||||
|
||||
if "</task>" not in content or "<details>" not in content or "</details>" not in content:
|
||||
return "", ""
|
||||
|
||||
desc_task, _, desc_details = content.rpartition("</task>")
|
||||
desc_details = desc_details.replace("<details>", "[italic]").replace(
|
||||
"</details>", "[/italic]"
|
||||
)
|
||||
return desc_task, desc_details
|
||||
|
||||
async def extract_tag_in_step(self, step: str) -> list[str]:
|
||||
steps_fmt = "\n\n".join(
|
||||
f'<step id="{ind + 1}">\n{s.strip()}\n</step>' for ind, s in enumerate(self.steps)
|
||||
)
|
||||
|
||||
if len(steps_fmt) > 300_000:
|
||||
# step_fmt is too long, skip tagging
|
||||
return []
|
||||
|
||||
llm_messages = [
|
||||
LLMMessage(
|
||||
role="user",
|
||||
content=f"Below is the trajectory of an AI agent solving a software bug until the current step. Each step is marked within a <step> tag.\n\n{steps_fmt}\n\n<current_step>{step}</current_step>",
|
||||
),
|
||||
LLMMessage(role="assistant", content="I understand."),
|
||||
LLMMessage(role="user", content=TAGGER_PROMPT),
|
||||
LLMMessage(role="assistant", content="Sure. The tags are: <tags>"),
|
||||
]
|
||||
self.model_config.temperature = 0.1
|
||||
|
||||
retry = 0
|
||||
while retry < 10:
|
||||
llm_response = self.lakeview_llm_client.chat(
|
||||
model_config=self.model_config,
|
||||
messages=llm_messages,
|
||||
reuse_history=False,
|
||||
)
|
||||
|
||||
content = "<tags>" + llm_response.content.lstrip()
|
||||
|
||||
matched_tags: list[str] = tags_re.findall(content)
|
||||
tags: list[str] = [tag.strip() for tag in matched_tags[0].split(",")]
|
||||
if all(tag in KNOWN_TAGS for tag in tags):
|
||||
return tags
|
||||
|
||||
retry += 1
|
||||
|
||||
return []
|
||||
|
||||
def _agent_step_str(self, agent_step: AgentStep) -> str | None:
|
||||
if agent_step.llm_response is None:
|
||||
return None
|
||||
|
||||
content = agent_step.llm_response.content.strip()
|
||||
|
||||
tool_calls_content = ""
|
||||
if agent_step.llm_response.tool_calls is not None:
|
||||
tool_calls_content = "\n".join(
|
||||
f"[`{tool_call.name}`] `{tool_call.arguments}`"
|
||||
for tool_call in agent_step.llm_response.tool_calls
|
||||
)
|
||||
tool_calls_content = tool_calls_content.strip()
|
||||
content = f"{content}\n\nTool calls:\n{tool_calls_content}"
|
||||
|
||||
return content
|
||||
|
||||
async def create_lakeview_step(self, agent_step: AgentStep) -> LakeViewStep | None:
|
||||
previous_step_str = "(none)"
|
||||
if len(self.steps) > 1:
|
||||
previous_step_str = self.steps[-1]
|
||||
|
||||
this_step_str = self._agent_step_str(agent_step)
|
||||
|
||||
if this_step_str:
|
||||
desc_task, desc_details = await self.extract_task_in_step(
|
||||
previous_step_str, this_step_str
|
||||
)
|
||||
tags = await self.extract_tag_in_step(this_step_str)
|
||||
tags_emoji = self.get_label(tags)
|
||||
return LakeViewStep(desc_task, desc_details, tags_emoji)
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,170 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
# TODO: remove these annotations by defining fine-grained types
|
||||
# pyright: reportAny=false
|
||||
# pyright: reportUnannotatedClassAttribute=false
|
||||
# pyright: reportUnknownMemberType=false
|
||||
# pyright: reportUnknownArgumentType=false
|
||||
# pyright: reportUnknownVariableType=false
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, override
|
||||
|
||||
|
||||
# data class for model parameters
|
||||
@dataclass
|
||||
class ModelParameters:
|
||||
"""Model parameters for a model provider."""
|
||||
|
||||
model: str
|
||||
api_key: str
|
||||
max_tokens: int
|
||||
temperature: float
|
||||
top_p: float
|
||||
top_k: int
|
||||
parallel_tool_calls: bool
|
||||
max_retries: int
|
||||
base_url: str | None = None
|
||||
api_version: str | None = None
|
||||
candidate_count: int | None = None # Gemini specific field
|
||||
stop_sequences: list[str] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class LakeviewConfig:
|
||||
"""Configuration for Lakeview."""
|
||||
|
||||
model_provider: str
|
||||
model_name: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class MCPServerConfig:
|
||||
# For stdio transport
|
||||
command: str | None = None
|
||||
args: list[str] | None = None
|
||||
env: dict[str, str] | None = None
|
||||
cwd: str | None = None
|
||||
|
||||
# For sse transport
|
||||
url: str | None = None
|
||||
|
||||
# For streamable http transport
|
||||
http_url: str | None = None
|
||||
headers: dict[str, str] | None = None
|
||||
|
||||
# For websocket transport
|
||||
tcp: str | None = None
|
||||
|
||||
# Common
|
||||
timeout: int | None = None
|
||||
trust: bool | None = None
|
||||
|
||||
# Metadata
|
||||
description: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class LegacyConfig:
|
||||
"""Configuration manager for Trae Agent."""
|
||||
|
||||
default_provider: str
|
||||
max_steps: int
|
||||
model_providers: dict[str, ModelParameters]
|
||||
mcp_servers: dict[str, MCPServerConfig]
|
||||
lakeview_config: LakeviewConfig | None = None
|
||||
enable_lakeview: bool = True
|
||||
allow_mcp_servers: list[str] = field(default_factory=list)
|
||||
|
||||
def __init__(self, config_or_config_file: str | dict = "trae_config.json"): # pyright: ignore[reportMissingTypeArgument, reportUnknownParameterType]
|
||||
# Accept either file path or direct config dict
|
||||
if isinstance(config_or_config_file, dict):
|
||||
self._config = config_or_config_file
|
||||
else:
|
||||
config_path = Path(config_or_config_file)
|
||||
if config_path.exists():
|
||||
try:
|
||||
with open(config_path, "r") as f:
|
||||
self._config = json.load(f)
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not load config file {config_or_config_file}: {e}")
|
||||
self._config = {}
|
||||
else:
|
||||
self._config = {}
|
||||
|
||||
self.default_provider = self._config.get("default_provider", "anthropic")
|
||||
self.max_steps = self._config.get("max_steps", 20)
|
||||
self.model_providers = {}
|
||||
self.enable_lakeview = self._config.get("enable_lakeview", True)
|
||||
self.mcp_servers = {
|
||||
k: MCPServerConfig(**v) for k, v in self._config.get("mcp_servers", {}).items()
|
||||
}
|
||||
self.allow_mcp_servers = self._config.get("allow_mcp_servers", [])
|
||||
|
||||
if len(self._config.get("model_providers", [])) == 0:
|
||||
self.model_providers = {
|
||||
"anthropic": ModelParameters(
|
||||
model="claude-sonnet-4-20250514",
|
||||
api_key="",
|
||||
base_url="https://api.anthropic.com",
|
||||
max_tokens=4096,
|
||||
temperature=0.5,
|
||||
top_p=1,
|
||||
top_k=0,
|
||||
parallel_tool_calls=False,
|
||||
max_retries=10,
|
||||
),
|
||||
}
|
||||
else:
|
||||
for provider in self._config.get("model_providers", {}):
|
||||
provider_config: dict[str, Any] = self._config.get("model_providers", {}).get(
|
||||
provider, {}
|
||||
)
|
||||
|
||||
candidate_count = provider_config.get("candidate_count")
|
||||
self.model_providers[provider] = ModelParameters(
|
||||
model=str(provider_config.get("model", "")),
|
||||
api_key=str(provider_config.get("api_key", "")),
|
||||
base_url=str(provider_config.get("base_url"))
|
||||
if "base_url" in provider_config
|
||||
else None,
|
||||
max_tokens=int(provider_config.get("max_tokens", 1000)),
|
||||
temperature=float(provider_config.get("temperature", 0.5)),
|
||||
top_p=float(provider_config.get("top_p", 1)),
|
||||
top_k=int(provider_config.get("top_k", 0)),
|
||||
max_retries=int(provider_config.get("max_retries", 10)),
|
||||
parallel_tool_calls=bool(provider_config.get("parallel_tool_calls", False)),
|
||||
api_version=str(provider_config.get("api_version"))
|
||||
if "api_version" in provider_config
|
||||
else None,
|
||||
candidate_count=int(candidate_count) if candidate_count is not None else None,
|
||||
stop_sequences=provider_config.get("stop_sequences")
|
||||
if "stop_sequences" in provider_config
|
||||
else None,
|
||||
)
|
||||
|
||||
# Configure lakeview_config - default to using default_provider settings
|
||||
lakeview_config_data = self._config.get("lakeview_config", {})
|
||||
if self.enable_lakeview:
|
||||
model_provider = lakeview_config_data.get("model_provider", None)
|
||||
model_name = lakeview_config_data.get("model_name", None)
|
||||
|
||||
if model_provider is None:
|
||||
model_provider = self.default_provider
|
||||
|
||||
if model_name is None:
|
||||
model_name = self.model_providers[model_provider].model
|
||||
|
||||
self.lakeview_config = LakeviewConfig(
|
||||
model_provider=str(model_provider),
|
||||
model_name=str(model_name),
|
||||
)
|
||||
|
||||
return
|
||||
|
||||
@override
|
||||
def __str__(self) -> str:
|
||||
return f"Config(default_provider={self.default_provider}, max_steps={self.max_steps}, model_providers={self.model_providers})"
|
||||
@@ -0,0 +1,220 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Anthropic API client wrapper with tool integration."""
|
||||
|
||||
import json
|
||||
from typing import override
|
||||
|
||||
import anthropic
|
||||
from anthropic.types.tool_union_param import TextEditor20250429
|
||||
|
||||
from trae_agent.tools.base import Tool, ToolCall, ToolResult
|
||||
from trae_agent.utils.config import ModelConfig
|
||||
from trae_agent.utils.llm_clients.base_client import BaseLLMClient
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage, LLMResponse, LLMUsage
|
||||
from trae_agent.utils.llm_clients.retry_utils import retry_with
|
||||
|
||||
|
||||
class AnthropicClient(BaseLLMClient):
|
||||
"""Anthropic client wrapper with tool schema generation."""
|
||||
|
||||
def __init__(self, model_config: ModelConfig):
|
||||
super().__init__(model_config)
|
||||
|
||||
self.client: anthropic.Anthropic = anthropic.Anthropic(
|
||||
api_key=self.api_key, base_url=self.base_url
|
||||
)
|
||||
self.message_history: list[anthropic.types.MessageParam] = []
|
||||
self.system_message: str | anthropic.NotGiven = anthropic.NOT_GIVEN
|
||||
|
||||
@override
|
||||
def set_chat_history(self, messages: list[LLMMessage]) -> None:
|
||||
"""Set the chat history."""
|
||||
self.message_history = self.parse_messages(messages)
|
||||
|
||||
def _create_anthropic_response(
|
||||
self,
|
||||
model_config: ModelConfig,
|
||||
tool_schemas: list[anthropic.types.ToolUnionParam] | anthropic.NotGiven,
|
||||
) -> anthropic.types.Message:
|
||||
"""Create a response using Anthropic API. This method will be decorated with retry logic."""
|
||||
return self.client.messages.create(
|
||||
model=model_config.model,
|
||||
messages=self.message_history,
|
||||
max_tokens=model_config.max_tokens,
|
||||
system=self.system_message,
|
||||
tools=tool_schemas,
|
||||
temperature=model_config.temperature,
|
||||
top_p=model_config.top_p,
|
||||
top_k=model_config.top_k,
|
||||
)
|
||||
|
||||
@override
|
||||
def chat(
|
||||
self,
|
||||
messages: list[LLMMessage],
|
||||
model_config: ModelConfig,
|
||||
tools: list[Tool] | None = None,
|
||||
reuse_history: bool = True,
|
||||
) -> LLMResponse:
|
||||
"""Send chat messages to Anthropic with optional tool support."""
|
||||
# Convert messages to Anthropic format
|
||||
anthropic_messages: list[anthropic.types.MessageParam] = self.parse_messages(messages)
|
||||
|
||||
self.message_history = (
|
||||
self.message_history + anthropic_messages if reuse_history else anthropic_messages
|
||||
)
|
||||
|
||||
# Add tools if provided
|
||||
tool_schemas: list[anthropic.types.ToolUnionParam] | anthropic.NotGiven = (
|
||||
anthropic.NOT_GIVEN
|
||||
)
|
||||
if tools:
|
||||
tool_schemas = []
|
||||
for tool in tools:
|
||||
if tool.name == "str_replace_based_edit_tool":
|
||||
tool_schemas.append(
|
||||
TextEditor20250429(
|
||||
name="str_replace_based_edit_tool",
|
||||
type="text_editor_20250429",
|
||||
)
|
||||
)
|
||||
elif tool.name == "bash":
|
||||
tool_schemas.append(
|
||||
anthropic.types.ToolBash20250124Param(name="bash", type="bash_20250124")
|
||||
)
|
||||
else:
|
||||
tool_schemas.append(
|
||||
anthropic.types.ToolParam(
|
||||
name=tool.name,
|
||||
description=tool.description,
|
||||
input_schema=tool.get_input_schema(),
|
||||
)
|
||||
)
|
||||
|
||||
# Apply retry decorator to the API call
|
||||
retry_decorator = retry_with(
|
||||
func=self._create_anthropic_response,
|
||||
provider_name="Anthropic",
|
||||
max_retries=model_config.max_retries,
|
||||
)
|
||||
response = retry_decorator(model_config, tool_schemas)
|
||||
|
||||
# Handle tool calls in response
|
||||
content = ""
|
||||
tool_calls: list[ToolCall] = []
|
||||
|
||||
for content_block in response.content:
|
||||
if content_block.type == "text":
|
||||
content += content_block.text
|
||||
self.message_history.append(
|
||||
anthropic.types.MessageParam(role="assistant", content=content_block.text)
|
||||
)
|
||||
elif content_block.type == "tool_use":
|
||||
tool_calls.append(
|
||||
ToolCall(
|
||||
call_id=content_block.id,
|
||||
name=content_block.name,
|
||||
arguments=content_block.input, # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
)
|
||||
self.message_history.append(
|
||||
anthropic.types.MessageParam(role="assistant", content=[content_block])
|
||||
)
|
||||
|
||||
usage = None
|
||||
if response.usage:
|
||||
usage = LLMUsage(
|
||||
input_tokens=response.usage.input_tokens or 0,
|
||||
output_tokens=response.usage.output_tokens or 0,
|
||||
cache_creation_input_tokens=response.usage.cache_creation_input_tokens or 0,
|
||||
cache_read_input_tokens=response.usage.cache_read_input_tokens or 0,
|
||||
)
|
||||
|
||||
llm_response = LLMResponse(
|
||||
content=content,
|
||||
usage=usage,
|
||||
model=response.model,
|
||||
finish_reason=response.stop_reason,
|
||||
tool_calls=tool_calls if len(tool_calls) > 0 else None,
|
||||
)
|
||||
|
||||
# Record trajectory if recorder is available
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_llm_interaction(
|
||||
messages=messages,
|
||||
response=llm_response,
|
||||
provider="anthropic",
|
||||
model=model_config.model,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
return llm_response
|
||||
|
||||
def parse_messages(self, messages: list[LLMMessage]) -> list[anthropic.types.MessageParam]:
|
||||
"""Parse the messages to Anthropic format."""
|
||||
anthropic_messages: list[anthropic.types.MessageParam] = []
|
||||
for msg in messages:
|
||||
if msg.role == "system":
|
||||
self.system_message = msg.content if msg.content else anthropic.NOT_GIVEN
|
||||
elif msg.tool_result:
|
||||
anthropic_messages.append(
|
||||
anthropic.types.MessageParam(
|
||||
role="user",
|
||||
content=[self.parse_tool_call_result(msg.tool_result)],
|
||||
)
|
||||
)
|
||||
elif msg.tool_call:
|
||||
anthropic_messages.append(
|
||||
anthropic.types.MessageParam(
|
||||
role="assistant", content=[self.parse_tool_call(msg.tool_call)]
|
||||
)
|
||||
)
|
||||
else:
|
||||
if msg.role == "user":
|
||||
role = "user"
|
||||
elif msg.role == "assistant":
|
||||
role = "assistant"
|
||||
else:
|
||||
raise ValueError(f"Invalid message role: {msg.role}")
|
||||
|
||||
if not msg.content:
|
||||
raise ValueError("Message content is required")
|
||||
|
||||
anthropic_messages.append(
|
||||
anthropic.types.MessageParam(role=role, content=msg.content)
|
||||
)
|
||||
return anthropic_messages
|
||||
|
||||
def parse_tool_call(self, tool_call: ToolCall) -> anthropic.types.ToolUseBlockParam:
|
||||
"""Parse the tool call from the LLM response."""
|
||||
return anthropic.types.ToolUseBlockParam(
|
||||
type="tool_use",
|
||||
id=tool_call.call_id,
|
||||
name=tool_call.name,
|
||||
input=json.dumps(tool_call.arguments),
|
||||
)
|
||||
|
||||
def parse_tool_call_result(
|
||||
self, tool_call_result: ToolResult
|
||||
) -> anthropic.types.ToolResultBlockParam:
|
||||
"""Parse the tool call result from the LLM response."""
|
||||
result: str = ""
|
||||
if tool_call_result.result:
|
||||
result = result + tool_call_result.result + "\n"
|
||||
if tool_call_result.error:
|
||||
result += "Tool call failed with error:\n"
|
||||
result += tool_call_result.error
|
||||
result = result.strip()
|
||||
|
||||
# Provide a default error message if the tool failed but didn't provide details
|
||||
if not tool_call_result.success and not result:
|
||||
result = "Tool execution failed without providing error details."
|
||||
|
||||
return anthropic.types.ToolResultBlockParam(
|
||||
tool_use_id=tool_call_result.call_id,
|
||||
type="tool_result",
|
||||
content=result,
|
||||
is_error=not tool_call_result.success,
|
||||
)
|
||||
@@ -0,0 +1,53 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Azure client wrapper with tool integrations"""
|
||||
|
||||
import openai
|
||||
|
||||
from trae_agent.utils.config import ModelConfig
|
||||
from trae_agent.utils.llm_clients.openai_compatible_base import (
|
||||
OpenAICompatibleClient,
|
||||
ProviderConfig,
|
||||
)
|
||||
|
||||
|
||||
class AzureProvider(ProviderConfig):
|
||||
"""Azure OpenAI provider configuration."""
|
||||
|
||||
def create_client(
|
||||
self, api_key: str, base_url: str | None, api_version: str | None
|
||||
) -> openai.OpenAI:
|
||||
"""Create Azure OpenAI client."""
|
||||
if not base_url:
|
||||
raise ValueError("base_url is required for AzureClient")
|
||||
|
||||
return openai.AzureOpenAI(
|
||||
azure_endpoint=base_url,
|
||||
api_version=api_version,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
def get_service_name(self) -> str:
|
||||
"""Get the service name for retry logging."""
|
||||
return "Azure OpenAI"
|
||||
|
||||
def get_provider_name(self) -> str:
|
||||
"""Get the provider name for trajectory recording."""
|
||||
return "azure"
|
||||
|
||||
def get_extra_headers(self) -> dict[str, str]:
|
||||
"""Get Azure-specific headers (none needed)."""
|
||||
return {}
|
||||
|
||||
def supports_tool_calling(self, model_name: str) -> bool:
|
||||
"""Check if the model supports tool calling."""
|
||||
# Azure OpenAI models generally support tool calling
|
||||
return True
|
||||
|
||||
|
||||
class AzureClient(OpenAICompatibleClient):
|
||||
"""Azure client wrapper that maintains compatibility while using the new architecture."""
|
||||
|
||||
def __init__(self, model_config: ModelConfig):
|
||||
super().__init__(model_config, AzureProvider())
|
||||
@@ -0,0 +1,44 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from trae_agent.tools.base import Tool
|
||||
from trae_agent.utils.config import ModelConfig
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage, LLMResponse
|
||||
from trae_agent.utils.trajectory_recorder import TrajectoryRecorder
|
||||
|
||||
|
||||
class BaseLLMClient(ABC):
|
||||
"""Base class for LLM clients."""
|
||||
|
||||
def __init__(self, model_config: ModelConfig):
|
||||
self.api_key: str = model_config.model_provider.api_key
|
||||
self.base_url: str | None = model_config.model_provider.base_url
|
||||
self.api_version: str | None = model_config.model_provider.api_version
|
||||
self.trajectory_recorder: TrajectoryRecorder | None = None # TrajectoryRecorder instance
|
||||
|
||||
def set_trajectory_recorder(self, recorder: TrajectoryRecorder | None) -> None:
|
||||
"""Set the trajectory recorder for this client."""
|
||||
self.trajectory_recorder = recorder
|
||||
|
||||
@abstractmethod
|
||||
def set_chat_history(self, messages: list[LLMMessage]) -> None:
|
||||
"""Set the chat history."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def chat(
|
||||
self,
|
||||
messages: list[LLMMessage],
|
||||
model_config: ModelConfig,
|
||||
tools: list[Tool] | None = None,
|
||||
reuse_history: bool = True,
|
||||
) -> LLMResponse:
|
||||
"""Send chat messages to the LLM."""
|
||||
pass
|
||||
|
||||
def supports_tool_calling(self, model_config: ModelConfig) -> bool:
|
||||
"""Check if the current model supports tool calling."""
|
||||
return model_config.supports_tool_calling
|
||||
@@ -0,0 +1,46 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Doubao client wrapper with tool integrations"""
|
||||
|
||||
import openai
|
||||
|
||||
from trae_agent.utils.config import ModelConfig
|
||||
from trae_agent.utils.llm_clients.openai_compatible_base import (
|
||||
OpenAICompatibleClient,
|
||||
ProviderConfig,
|
||||
)
|
||||
|
||||
|
||||
class DoubaoProvider(ProviderConfig):
|
||||
"""Doubao provider configuration."""
|
||||
|
||||
def create_client(
|
||||
self, api_key: str, base_url: str | None, api_version: str | None
|
||||
) -> openai.OpenAI:
|
||||
"""Create OpenAI client with Doubao base URL."""
|
||||
return openai.OpenAI(base_url=base_url, api_key=api_key)
|
||||
|
||||
def get_service_name(self) -> str:
|
||||
"""Get the service name for retry logging."""
|
||||
return "Doubao"
|
||||
|
||||
def get_provider_name(self) -> str:
|
||||
"""Get the provider name for trajectory recording."""
|
||||
return "doubao"
|
||||
|
||||
def get_extra_headers(self) -> dict[str, str]:
|
||||
"""Get Doubao-specific headers (none needed)."""
|
||||
return {}
|
||||
|
||||
def supports_tool_calling(self, model_name: str) -> bool:
|
||||
"""Check if the model supports tool calling."""
|
||||
# Doubao models generally support tool calling
|
||||
return True
|
||||
|
||||
|
||||
class DoubaoClient(OpenAICompatibleClient):
|
||||
"""Doubao client wrapper that maintains compatibility while using the new architecture."""
|
||||
|
||||
def __init__(self, model_config: ModelConfig):
|
||||
super().__init__(model_config, DoubaoProvider())
|
||||
@@ -0,0 +1,229 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Google Gemini API client wrapper with tool integration."""
|
||||
|
||||
import json
|
||||
import traceback
|
||||
import uuid
|
||||
from typing import override
|
||||
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
|
||||
from trae_agent.tools.base import Tool, ToolCall, ToolResult
|
||||
from trae_agent.utils.config import ModelConfig
|
||||
from trae_agent.utils.llm_clients.base_client import BaseLLMClient
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage, LLMResponse, LLMUsage
|
||||
from trae_agent.utils.llm_clients.retry_utils import retry_with
|
||||
|
||||
|
||||
class GoogleClient(BaseLLMClient):
|
||||
"""Google Gemini client wrapper with tool schema generation."""
|
||||
|
||||
def __init__(self, model_config: ModelConfig):
|
||||
super().__init__(model_config)
|
||||
|
||||
self.client = genai.Client(api_key=self.api_key)
|
||||
self.message_history: list[types.Content] = []
|
||||
self.system_instruction: str | None = None
|
||||
|
||||
@override
|
||||
def set_chat_history(self, messages: list[LLMMessage]) -> None:
|
||||
"""Set the chat history."""
|
||||
self.message_history, self.system_instruction = self.parse_messages(messages)
|
||||
|
||||
def _create_google_response(
|
||||
self,
|
||||
model_config: ModelConfig,
|
||||
current_chat_contents: list[types.Content],
|
||||
generation_config: types.GenerateContentConfig,
|
||||
) -> types.GenerateContentResponse:
|
||||
"""Create a response using Google Gemini API. This method will be decorated with retry logic."""
|
||||
return self.client.models.generate_content( # pyright: ignore[reportUnknownMemberType]
|
||||
model=model_config.model,
|
||||
contents=current_chat_contents,
|
||||
config=generation_config,
|
||||
)
|
||||
|
||||
@override
|
||||
def chat(
|
||||
self,
|
||||
messages: list[LLMMessage],
|
||||
model_config: ModelConfig,
|
||||
tools: list[Tool] | None = None,
|
||||
reuse_history: bool = True,
|
||||
) -> LLMResponse:
|
||||
"""Send chat messages to Gemini with optional tool support."""
|
||||
newly_parsed_messages, system_instruction_from_message = self.parse_messages(messages)
|
||||
|
||||
current_system_instruction = system_instruction_from_message or self.system_instruction
|
||||
|
||||
if reuse_history:
|
||||
current_chat_contents = self.message_history + newly_parsed_messages
|
||||
else:
|
||||
current_chat_contents = newly_parsed_messages
|
||||
|
||||
# Set up generation config
|
||||
generation_config = types.GenerateContentConfig(
|
||||
temperature=model_config.temperature,
|
||||
top_p=model_config.top_p,
|
||||
top_k=model_config.top_k,
|
||||
max_output_tokens=model_config.max_tokens,
|
||||
candidate_count=model_config.candidate_count,
|
||||
stop_sequences=model_config.stop_sequences,
|
||||
system_instruction=current_system_instruction,
|
||||
)
|
||||
|
||||
# Add tools if provided
|
||||
if tools:
|
||||
tool_schemas = [
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(
|
||||
name=tool.get_name(),
|
||||
description=tool.get_description(),
|
||||
parameters=tool.get_input_schema(), # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
]
|
||||
)
|
||||
for tool in tools
|
||||
]
|
||||
generation_config.tools = tool_schemas
|
||||
|
||||
# Apply retry decorator to the API call
|
||||
retry_decorator = retry_with(
|
||||
func=self._create_google_response,
|
||||
provider_name="Google Gemini",
|
||||
max_retries=model_config.max_retries,
|
||||
)
|
||||
response = retry_decorator(model_config, current_chat_contents, generation_config)
|
||||
|
||||
content = ""
|
||||
tool_calls: list[ToolCall] = []
|
||||
assistant_response_content = None
|
||||
|
||||
if response.candidates:
|
||||
candidate = response.candidates[0]
|
||||
if candidate.content and candidate.content.parts:
|
||||
assistant_response_content = candidate.content
|
||||
for part in candidate.content.parts:
|
||||
if part.text:
|
||||
content += part.text
|
||||
elif part.function_call:
|
||||
tool_calls.append(
|
||||
ToolCall(
|
||||
call_id=str(uuid.uuid4()),
|
||||
name=part.function_call.name or "tool",
|
||||
arguments=dict(part.function_call.args)
|
||||
if part.function_call.args
|
||||
else {},
|
||||
)
|
||||
)
|
||||
|
||||
if reuse_history:
|
||||
new_history = self.message_history + newly_parsed_messages
|
||||
else:
|
||||
new_history = newly_parsed_messages
|
||||
|
||||
if assistant_response_content:
|
||||
new_history.append(assistant_response_content)
|
||||
|
||||
self.message_history = new_history
|
||||
|
||||
if current_system_instruction:
|
||||
self.system_instruction = current_system_instruction
|
||||
|
||||
usage = None
|
||||
if response.usage_metadata:
|
||||
usage = LLMUsage(
|
||||
input_tokens=response.usage_metadata.prompt_token_count or 0,
|
||||
output_tokens=response.usage_metadata.candidates_token_count or 0,
|
||||
cache_read_input_tokens=response.usage_metadata.cached_content_token_count or 0,
|
||||
cache_creation_input_tokens=0,
|
||||
)
|
||||
|
||||
llm_response = LLMResponse(
|
||||
content=content,
|
||||
usage=usage,
|
||||
model=model_config.model,
|
||||
finish_reason=str(
|
||||
response.candidates[0].finish_reason.name
|
||||
if response.candidates[0].finish_reason
|
||||
else "unknown"
|
||||
)
|
||||
if response.candidates
|
||||
else "UNKNOWN",
|
||||
tool_calls=tool_calls if len(tool_calls) > 0 else None,
|
||||
)
|
||||
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_llm_interaction(
|
||||
messages=messages,
|
||||
response=llm_response,
|
||||
provider="google",
|
||||
model=model_config.model,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
return llm_response
|
||||
|
||||
def parse_messages(self, messages: list[LLMMessage]) -> tuple[list[types.Content], str | None]:
|
||||
"""Parse the messages to Gemini format, separating system instructions."""
|
||||
gemini_messages: list[types.Content] = []
|
||||
system_instruction: str | None = None
|
||||
|
||||
for msg in messages:
|
||||
if msg.role == "system":
|
||||
system_instruction = msg.content
|
||||
continue
|
||||
elif msg.tool_result:
|
||||
gemini_messages.append(
|
||||
types.Content(
|
||||
role="tool",
|
||||
parts=[self.parse_tool_call_result(msg.tool_result)],
|
||||
)
|
||||
)
|
||||
elif msg.tool_call:
|
||||
gemini_messages.append(
|
||||
types.Content(role="model", parts=[self.parse_tool_call(msg.tool_call)])
|
||||
)
|
||||
else:
|
||||
role = "user" if msg.role == "user" else "model"
|
||||
gemini_messages.append(
|
||||
types.Content(role=role, parts=[types.Part(text=msg.content or "")])
|
||||
)
|
||||
|
||||
return gemini_messages, system_instruction
|
||||
|
||||
def parse_tool_call(self, tool_call: ToolCall) -> types.Part:
|
||||
"""Parse a ToolCall into a Gemini FunctionCall Part for history."""
|
||||
return types.Part.from_function_call(name=tool_call.name, args=tool_call.arguments)
|
||||
|
||||
def parse_tool_call_result(self, tool_result: ToolResult) -> types.Part:
|
||||
"""Parse a ToolResult into a Gemini FunctionResponse Part for history."""
|
||||
result_content: dict[str, str] = {}
|
||||
if tool_result.result is not None:
|
||||
try:
|
||||
json.dumps(tool_result.result)
|
||||
result_content["result"] = tool_result.result
|
||||
except (TypeError, OverflowError) as e:
|
||||
tb = traceback.format_exc()
|
||||
serialization_error = f"JSON serialization failed for tool result: {e}\n{tb}"
|
||||
if tool_result.error:
|
||||
result_content["error"] = f"{tool_result.error}\n\n{serialization_error}"
|
||||
else:
|
||||
result_content["error"] = serialization_error
|
||||
result_content["result"] = str(tool_result.result)
|
||||
|
||||
if tool_result.error and "error" not in result_content:
|
||||
result_content["error"] = tool_result.error
|
||||
|
||||
if not result_content:
|
||||
result_content["status"] = "Tool executed successfully but returned no output."
|
||||
|
||||
if not hasattr(tool_result, "name") or not tool_result.name:
|
||||
raise AttributeError(
|
||||
"ToolResult must have a 'name' attribute matching the function that was called."
|
||||
)
|
||||
return types.Part.from_function_response(name=tool_result.name, response=result_content)
|
||||
@@ -0,0 +1,52 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from trae_agent.tools.base import ToolCall, ToolResult
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMMessage:
|
||||
"""Standard message format."""
|
||||
|
||||
role: str
|
||||
content: str | None = None
|
||||
tool_call: ToolCall | None = None
|
||||
tool_result: ToolResult | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMUsage:
|
||||
"""LLM usage format."""
|
||||
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
cache_creation_input_tokens: int = 0
|
||||
cache_read_input_tokens: int = 0
|
||||
reasoning_tokens: int = 0
|
||||
|
||||
def __add__(self, other: "LLMUsage") -> "LLMUsage":
|
||||
return LLMUsage(
|
||||
input_tokens=self.input_tokens + other.input_tokens,
|
||||
output_tokens=self.output_tokens + other.output_tokens,
|
||||
cache_creation_input_tokens=self.cache_creation_input_tokens
|
||||
+ other.cache_creation_input_tokens,
|
||||
cache_read_input_tokens=self.cache_read_input_tokens + other.cache_read_input_tokens,
|
||||
reasoning_tokens=self.reasoning_tokens + other.reasoning_tokens,
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"LLMUsage(input_tokens={self.input_tokens}, output_tokens={self.output_tokens}, cache_creation_input_tokens={self.cache_creation_input_tokens}, cache_read_input_tokens={self.cache_read_input_tokens}, reasoning_tokens={self.reasoning_tokens})"
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMResponse:
|
||||
"""Standard LLM response format."""
|
||||
|
||||
content: str
|
||||
usage: LLMUsage | None = None
|
||||
model: str | None = None
|
||||
finish_reason: str | None = None
|
||||
tool_calls: list[ToolCall] | None = None
|
||||
@@ -0,0 +1,86 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""LLM Client wrapper for OpenAI, Anthropic, Azure, and OpenRouter APIs."""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
from trae_agent.tools.base import Tool
|
||||
from trae_agent.utils.config import ModelConfig
|
||||
from trae_agent.utils.llm_clients.base_client import BaseLLMClient
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage, LLMResponse
|
||||
from trae_agent.utils.trajectory_recorder import TrajectoryRecorder
|
||||
|
||||
|
||||
class LLMProvider(Enum):
|
||||
"""Supported LLM providers."""
|
||||
|
||||
OPENAI = "openai"
|
||||
ANTHROPIC = "anthropic"
|
||||
AZURE = "azure"
|
||||
OLLAMA = "ollama"
|
||||
OPENROUTER = "openrouter"
|
||||
DOUBAO = "doubao"
|
||||
GOOGLE = "google"
|
||||
|
||||
|
||||
class LLMClient:
|
||||
"""Main LLM client that supports multiple providers."""
|
||||
|
||||
def __init__(self, model_config: ModelConfig):
|
||||
self.provider: LLMProvider = LLMProvider(model_config.model_provider.provider)
|
||||
self.model_config: ModelConfig = model_config
|
||||
|
||||
match self.provider:
|
||||
case LLMProvider.OPENAI:
|
||||
from .openai_client import OpenAIClient
|
||||
|
||||
self.client: BaseLLMClient = OpenAIClient(model_config)
|
||||
case LLMProvider.ANTHROPIC:
|
||||
from .anthropic_client import AnthropicClient
|
||||
|
||||
self.client = AnthropicClient(model_config)
|
||||
case LLMProvider.AZURE:
|
||||
from .azure_client import AzureClient
|
||||
|
||||
self.client = AzureClient(model_config)
|
||||
case LLMProvider.OPENROUTER:
|
||||
from .openrouter_client import OpenRouterClient
|
||||
|
||||
self.client = OpenRouterClient(model_config)
|
||||
case LLMProvider.DOUBAO:
|
||||
from .doubao_client import DoubaoClient
|
||||
|
||||
self.client = DoubaoClient(model_config)
|
||||
case LLMProvider.OLLAMA:
|
||||
from .ollama_client import OllamaClient
|
||||
|
||||
self.client = OllamaClient(model_config)
|
||||
case LLMProvider.GOOGLE:
|
||||
from .google_client import GoogleClient
|
||||
|
||||
self.client = GoogleClient(model_config)
|
||||
|
||||
def set_trajectory_recorder(self, recorder: TrajectoryRecorder | None) -> None:
|
||||
"""Set the trajectory recorder for the underlying client."""
|
||||
self.client.set_trajectory_recorder(recorder)
|
||||
|
||||
def set_chat_history(self, messages: list[LLMMessage]) -> None:
|
||||
"""Set the chat history."""
|
||||
self.client.set_chat_history(messages)
|
||||
|
||||
def chat(
|
||||
self,
|
||||
messages: list[LLMMessage],
|
||||
model_config: ModelConfig,
|
||||
tools: list[Tool] | None = None,
|
||||
reuse_history: bool = True,
|
||||
) -> LLMResponse:
|
||||
"""Send chat messages to the LLM."""
|
||||
return self.client.chat(messages, model_config, tools, reuse_history)
|
||||
|
||||
def supports_tool_calling(self, model_config: ModelConfig) -> bool:
|
||||
"""Check if the current client supports tool calling."""
|
||||
return hasattr(self.client, "supports_tool_calling") and self.client.supports_tool_calling(
|
||||
model_config
|
||||
)
|
||||
@@ -0,0 +1,196 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""
|
||||
Ollama API client wrapper with tool integration
|
||||
"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from typing import override
|
||||
|
||||
import openai
|
||||
from ollama import chat as ollama_chat # pyright: ignore[reportUnknownVariableType]
|
||||
from openai.types.responses import (
|
||||
FunctionToolParam,
|
||||
ResponseFunctionToolCallParam,
|
||||
ResponseInputParam,
|
||||
)
|
||||
from openai.types.responses.response_input_param import FunctionCallOutput
|
||||
|
||||
from trae_agent.tools.base import Tool, ToolCall, ToolResult
|
||||
from trae_agent.utils.config import ModelConfig
|
||||
from trae_agent.utils.llm_clients.base_client import BaseLLMClient
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage, LLMResponse
|
||||
from trae_agent.utils.llm_clients.retry_utils import retry_with
|
||||
|
||||
|
||||
class OllamaClient(BaseLLMClient):
|
||||
def __init__(self, model_config: ModelConfig):
|
||||
super().__init__(model_config)
|
||||
|
||||
self.client: openai.OpenAI = openai.OpenAI(
|
||||
# by default ollama doesn't require any api key. It should set to be "ollama".
|
||||
api_key=self.api_key,
|
||||
base_url=model_config.model_provider.base_url
|
||||
if model_config.model_provider.base_url
|
||||
else "http://localhost:11434/v1",
|
||||
)
|
||||
|
||||
self.message_history: ResponseInputParam = []
|
||||
|
||||
@override
|
||||
def set_chat_history(self, messages: list[LLMMessage]) -> None:
|
||||
self.message_history = self.parse_messages(messages)
|
||||
|
||||
def _create_ollama_response(
|
||||
self,
|
||||
model_config: ModelConfig,
|
||||
tool_schemas: list[FunctionToolParam] | None,
|
||||
):
|
||||
"""Create a response using Ollama API. This method will be decorated with retry logic."""
|
||||
tools_param = None
|
||||
if tool_schemas:
|
||||
tools_param = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool["name"],
|
||||
"description": tool.get("description", ""),
|
||||
"parameters": tool["parameters"],
|
||||
},
|
||||
}
|
||||
for tool in tool_schemas
|
||||
]
|
||||
return ollama_chat(
|
||||
messages=self.message_history,
|
||||
model=model_config.model,
|
||||
tools=tools_param,
|
||||
)
|
||||
|
||||
@override
|
||||
def chat(
|
||||
self,
|
||||
messages: list[LLMMessage],
|
||||
model_config: ModelConfig,
|
||||
tools: list[Tool] | None = None,
|
||||
reuse_history: bool = True,
|
||||
) -> LLMResponse:
|
||||
"""
|
||||
A rewritten version of ollama chan
|
||||
"""
|
||||
msgs: ResponseInputParam = self.parse_messages(messages)
|
||||
|
||||
tool_schemas = None
|
||||
if tools:
|
||||
tool_schemas = [
|
||||
FunctionToolParam(
|
||||
name=tool.name,
|
||||
description=tool.description,
|
||||
parameters=tool.get_input_schema(),
|
||||
strict=True,
|
||||
type="function",
|
||||
)
|
||||
for tool in tools
|
||||
]
|
||||
|
||||
if reuse_history:
|
||||
self.message_history = self.message_history + msgs
|
||||
else:
|
||||
self.message_history = msgs
|
||||
|
||||
# Apply retry decorator to the API call
|
||||
retry_decorator = retry_with(
|
||||
func=self._create_ollama_response,
|
||||
provider_name="Ollama",
|
||||
max_retries=model_config.max_retries,
|
||||
)
|
||||
response = retry_decorator(model_config, tool_schemas)
|
||||
|
||||
content = ""
|
||||
tool_calls: list[ToolCall] = []
|
||||
|
||||
if response.message.tool_calls:
|
||||
for tool in response.message.tool_calls:
|
||||
tool_calls.append(
|
||||
ToolCall(
|
||||
call_id=self._id_generator(),
|
||||
name=tool.function.name,
|
||||
arguments=dict(tool.function.arguments),
|
||||
id=self._id_generator(),
|
||||
)
|
||||
)
|
||||
else:
|
||||
# consider response is not a tool call
|
||||
content = str(response.message.content)
|
||||
|
||||
llm_response = LLMResponse(
|
||||
content=content,
|
||||
usage=None,
|
||||
model=model_config.model,
|
||||
finish_reason=None, # seems can't get finish reason will check docs soon
|
||||
tool_calls=tool_calls if len(tool_calls) > 0 else None,
|
||||
)
|
||||
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_llm_interaction(
|
||||
messages=messages,
|
||||
response=llm_response,
|
||||
provider="ollama",
|
||||
model=model_config.model,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
return llm_response
|
||||
|
||||
def parse_messages(self, messages: list[LLMMessage]) -> ResponseInputParam:
|
||||
"""
|
||||
Ollama parse messages should be compatible with openai handling
|
||||
"""
|
||||
openai_messages: ResponseInputParam = []
|
||||
for msg in messages:
|
||||
if msg.tool_result:
|
||||
openai_messages.append(self.parse_tool_call_result(msg.tool_result))
|
||||
elif msg.tool_call:
|
||||
openai_messages.append(self.parse_tool_call(msg.tool_call))
|
||||
else:
|
||||
if not msg.content:
|
||||
raise ValueError("Message content is required")
|
||||
if msg.role == "system":
|
||||
openai_messages.append({"role": "system", "content": msg.content})
|
||||
elif msg.role == "user":
|
||||
openai_messages.append({"role": "user", "content": msg.content})
|
||||
elif msg.role == "assistant":
|
||||
openai_messages.append({"role": "assistant", "content": msg.content})
|
||||
else:
|
||||
raise ValueError(f"Invalid message role: {msg.role}")
|
||||
return openai_messages
|
||||
|
||||
def parse_tool_call(self, tool_call: ToolCall) -> ResponseFunctionToolCallParam:
|
||||
"""Parse the tool call from the LLM response."""
|
||||
return ResponseFunctionToolCallParam(
|
||||
call_id=tool_call.call_id,
|
||||
name=tool_call.name,
|
||||
arguments=json.dumps(tool_call.arguments),
|
||||
type="function_call",
|
||||
)
|
||||
|
||||
def parse_tool_call_result(self, tool_call_result: ToolResult) -> FunctionCallOutput:
|
||||
"""Parse the tool call result from the LLM response."""
|
||||
result: str = ""
|
||||
if tool_call_result.result:
|
||||
result = result + tool_call_result.result + "\n"
|
||||
if tool_call_result.error:
|
||||
result += tool_call_result.error
|
||||
result = result.strip()
|
||||
|
||||
return FunctionCallOutput(
|
||||
call_id=tool_call_result.call_id,
|
||||
id=tool_call_result.id,
|
||||
output=result,
|
||||
type="function_call_output",
|
||||
)
|
||||
|
||||
def _id_generator(self) -> str:
|
||||
"""Generate a random ID string"""
|
||||
return str(uuid.uuid4())
|
||||
@@ -0,0 +1,209 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""OpenAI API client wrapper with tool integration."""
|
||||
|
||||
import json
|
||||
from typing import override
|
||||
|
||||
import openai
|
||||
from openai.types.responses import (
|
||||
EasyInputMessageParam,
|
||||
FunctionToolParam,
|
||||
Response,
|
||||
ResponseFunctionToolCallParam,
|
||||
ResponseInputParam,
|
||||
ToolParam,
|
||||
)
|
||||
from openai.types.responses.response_input_param import FunctionCallOutput
|
||||
|
||||
from trae_agent.tools.base import Tool, ToolCall, ToolResult
|
||||
from trae_agent.utils.config import ModelConfig
|
||||
from trae_agent.utils.llm_clients.base_client import BaseLLMClient
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage, LLMResponse, LLMUsage
|
||||
from trae_agent.utils.llm_clients.retry_utils import retry_with
|
||||
|
||||
|
||||
class OpenAIClient(BaseLLMClient):
|
||||
"""OpenAI client wrapper with tool schema generation."""
|
||||
|
||||
def __init__(self, model_config: ModelConfig):
|
||||
super().__init__(model_config)
|
||||
|
||||
self.client: openai.OpenAI = openai.OpenAI(api_key=self.api_key, base_url=self.base_url)
|
||||
self.message_history: ResponseInputParam = []
|
||||
|
||||
@override
|
||||
def set_chat_history(self, messages: list[LLMMessage]) -> None:
|
||||
"""Set the chat history."""
|
||||
self.message_history = self.parse_messages(messages)
|
||||
|
||||
def _create_openai_response(
|
||||
self,
|
||||
api_call_input: ResponseInputParam,
|
||||
model_config: ModelConfig,
|
||||
tool_schemas: list[ToolParam] | None,
|
||||
) -> Response:
|
||||
"""Create a response using OpenAI API. This method will be decorated with retry logic."""
|
||||
return self.client.responses.create(
|
||||
input=api_call_input,
|
||||
model=model_config.model,
|
||||
tools=tool_schemas if tool_schemas else openai.NOT_GIVEN,
|
||||
temperature=model_config.temperature
|
||||
if "o3" not in model_config.model
|
||||
and "o4-mini" not in model_config.model
|
||||
and "gpt-5" not in model_config.model
|
||||
else openai.NOT_GIVEN,
|
||||
top_p=model_config.top_p,
|
||||
max_output_tokens=model_config.max_tokens,
|
||||
)
|
||||
|
||||
@override
|
||||
def chat(
|
||||
self,
|
||||
messages: list[LLMMessage],
|
||||
model_config: ModelConfig,
|
||||
tools: list[Tool] | None = None,
|
||||
reuse_history: bool = True,
|
||||
) -> LLMResponse:
|
||||
"""Send chat messages to OpenAI with optional tool support."""
|
||||
openai_messages: ResponseInputParam = self.parse_messages(messages)
|
||||
|
||||
if reuse_history:
|
||||
self.message_history = self.message_history + openai_messages
|
||||
else:
|
||||
self.message_history = openai_messages
|
||||
|
||||
tool_schemas = None
|
||||
if tools:
|
||||
tool_schemas = [
|
||||
FunctionToolParam(
|
||||
name=tool.name,
|
||||
description=tool.description,
|
||||
parameters=tool.get_input_schema(),
|
||||
strict=True,
|
||||
type="function",
|
||||
)
|
||||
for tool in tools
|
||||
]
|
||||
|
||||
api_call_input: ResponseInputParam = self.message_history
|
||||
|
||||
# Apply retry decorator to the API call
|
||||
retry_decorator = retry_with(
|
||||
func=self._create_openai_response,
|
||||
provider_name="OpenAI",
|
||||
max_retries=model_config.max_retries,
|
||||
)
|
||||
response = retry_decorator(api_call_input, model_config, tool_schemas)
|
||||
|
||||
content = ""
|
||||
tool_calls: list[ToolCall] = []
|
||||
for output_block in response.output:
|
||||
if output_block.type == "function_call":
|
||||
tool_calls.append(
|
||||
ToolCall(
|
||||
call_id=output_block.call_id,
|
||||
name=output_block.name,
|
||||
arguments=json.loads(output_block.arguments)
|
||||
if output_block.arguments
|
||||
else {},
|
||||
id=output_block.id,
|
||||
)
|
||||
)
|
||||
tool_call_param = ResponseFunctionToolCallParam(
|
||||
arguments=output_block.arguments,
|
||||
call_id=output_block.call_id,
|
||||
name=output_block.name,
|
||||
type="function_call",
|
||||
)
|
||||
if output_block.status:
|
||||
tool_call_param["status"] = output_block.status
|
||||
if output_block.id:
|
||||
tool_call_param["id"] = output_block.id
|
||||
self.message_history.append(tool_call_param)
|
||||
elif output_block.type == "message":
|
||||
content = "".join(
|
||||
content_block.text
|
||||
for content_block in output_block.content
|
||||
if content_block.type == "output_text"
|
||||
)
|
||||
|
||||
if content != "":
|
||||
self.message_history.append(
|
||||
EasyInputMessageParam(content=content, role="assistant", type="message")
|
||||
)
|
||||
|
||||
usage = None
|
||||
if response.usage:
|
||||
usage = LLMUsage(
|
||||
input_tokens=response.usage.input_tokens or 0,
|
||||
output_tokens=response.usage.output_tokens or 0,
|
||||
cache_read_input_tokens=response.usage.input_tokens_details.cached_tokens or 0,
|
||||
reasoning_tokens=response.usage.output_tokens_details.reasoning_tokens or 0,
|
||||
)
|
||||
|
||||
llm_response = LLMResponse(
|
||||
content=content,
|
||||
usage=usage,
|
||||
model=response.model,
|
||||
finish_reason=response.status,
|
||||
tool_calls=tool_calls if len(tool_calls) > 0 else None,
|
||||
)
|
||||
|
||||
# Record trajectory if recorder is available
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_llm_interaction(
|
||||
messages=messages,
|
||||
response=llm_response,
|
||||
provider="openai",
|
||||
model=model_config.model,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
return llm_response
|
||||
|
||||
def parse_messages(self, messages: list[LLMMessage]) -> ResponseInputParam:
|
||||
"""Parse the messages to OpenAI format."""
|
||||
openai_messages: ResponseInputParam = []
|
||||
for msg in messages:
|
||||
if msg.tool_result:
|
||||
openai_messages.append(self.parse_tool_call_result(msg.tool_result))
|
||||
elif msg.tool_call:
|
||||
openai_messages.append(self.parse_tool_call(msg.tool_call))
|
||||
else:
|
||||
if not msg.content:
|
||||
raise ValueError("Message content is required")
|
||||
if msg.role == "system":
|
||||
openai_messages.append({"role": "system", "content": msg.content})
|
||||
elif msg.role == "user":
|
||||
openai_messages.append({"role": "user", "content": msg.content})
|
||||
elif msg.role == "assistant":
|
||||
openai_messages.append({"role": "assistant", "content": msg.content})
|
||||
else:
|
||||
raise ValueError(f"Invalid message role: {msg.role}")
|
||||
return openai_messages
|
||||
|
||||
def parse_tool_call(self, tool_call: ToolCall) -> ResponseFunctionToolCallParam:
|
||||
"""Parse the tool call from the LLM response."""
|
||||
return ResponseFunctionToolCallParam(
|
||||
call_id=tool_call.call_id,
|
||||
name=tool_call.name,
|
||||
arguments=json.dumps(tool_call.arguments),
|
||||
type="function_call",
|
||||
)
|
||||
|
||||
def parse_tool_call_result(self, tool_call_result: ToolResult) -> FunctionCallOutput:
|
||||
"""Parse the tool call result from the LLM response to FunctionCallOutput format."""
|
||||
result_content: str = ""
|
||||
if tool_call_result.result is not None:
|
||||
result_content += str(tool_call_result.result)
|
||||
if tool_call_result.error:
|
||||
result_content += f"\nError: {tool_call_result.error}"
|
||||
result_content = result_content.strip()
|
||||
|
||||
return FunctionCallOutput(
|
||||
type="function_call_output", # Explicitly set the type field
|
||||
call_id=tool_call_result.call_id,
|
||||
output=result_content,
|
||||
)
|
||||
@@ -0,0 +1,286 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Base class for OpenAI-compatible clients with shared logic."""
|
||||
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import override
|
||||
|
||||
import openai
|
||||
from openai.types.chat import (
|
||||
ChatCompletion,
|
||||
ChatCompletionAssistantMessageParam,
|
||||
ChatCompletionFunctionMessageParam,
|
||||
ChatCompletionMessageParam,
|
||||
ChatCompletionMessageToolCallParam,
|
||||
ChatCompletionSystemMessageParam,
|
||||
ChatCompletionToolParam,
|
||||
ChatCompletionUserMessageParam,
|
||||
)
|
||||
from openai.types.chat.chat_completion_message_tool_call_param import Function
|
||||
from openai.types.chat.chat_completion_tool_message_param import (
|
||||
ChatCompletionToolMessageParam,
|
||||
)
|
||||
from openai.types.shared_params.function_definition import FunctionDefinition
|
||||
|
||||
from trae_agent.tools.base import Tool, ToolCall
|
||||
from trae_agent.utils.config import ModelConfig
|
||||
from trae_agent.utils.llm_clients.base_client import BaseLLMClient
|
||||
from trae_agent.utils.llm_clients.llm_basics import LLMMessage, LLMResponse, LLMUsage
|
||||
from trae_agent.utils.llm_clients.retry_utils import retry_with
|
||||
|
||||
|
||||
class ProviderConfig(ABC):
|
||||
"""Abstract base class for provider-specific configurations."""
|
||||
|
||||
@abstractmethod
|
||||
def create_client(
|
||||
self, api_key: str, base_url: str | None, api_version: str | None
|
||||
) -> openai.OpenAI:
|
||||
"""Create the OpenAI client instance."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_service_name(self) -> str:
|
||||
"""Get the service name for retry logging."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_provider_name(self) -> str:
|
||||
"""Get the provider name for trajectory recording."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_extra_headers(self) -> dict[str, str]:
|
||||
"""Get any extra headers needed for the API call."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def supports_tool_calling(self, model_name: str) -> bool:
|
||||
"""Check if the model supports tool calling."""
|
||||
pass
|
||||
|
||||
|
||||
class OpenAICompatibleClient(BaseLLMClient):
|
||||
"""Base class for OpenAI-compatible clients with shared logic."""
|
||||
|
||||
def __init__(self, model_config: ModelConfig, provider_config: ProviderConfig):
|
||||
super().__init__(model_config)
|
||||
self.provider_config = provider_config
|
||||
self.client = provider_config.create_client(self.api_key, self.base_url, self.api_version)
|
||||
self.message_history: list[ChatCompletionMessageParam] = []
|
||||
|
||||
@override
|
||||
def set_chat_history(self, messages: list[LLMMessage]) -> None:
|
||||
"""Set the chat history."""
|
||||
self.message_history = self.parse_messages(messages)
|
||||
|
||||
def _create_response(
|
||||
self,
|
||||
model_config: ModelConfig,
|
||||
tool_schemas: list[ChatCompletionToolParam] | None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
) -> ChatCompletion:
|
||||
"""Create a response using the provider's API. This method will be decorated with retry logic."""
|
||||
"""Select the correct token parameter based on model configuration.
|
||||
If max_completion_tokens is set, use it. Otherwise, use max_tokens."""
|
||||
token_params = {}
|
||||
if model_config.should_use_max_completion_tokens():
|
||||
token_params["max_completion_tokens"] = model_config.get_max_tokens_param()
|
||||
else:
|
||||
token_params["max_tokens"] = model_config.get_max_tokens_param()
|
||||
|
||||
return self.client.chat.completions.create(
|
||||
model=model_config.model,
|
||||
messages=self.message_history,
|
||||
tools=tool_schemas if tool_schemas else openai.NOT_GIVEN,
|
||||
temperature=model_config.temperature
|
||||
if "o3" not in model_config.model
|
||||
and "o4-mini" not in model_config.model
|
||||
and "gpt-5" not in model_config.model
|
||||
else openai.NOT_GIVEN,
|
||||
top_p=model_config.top_p,
|
||||
extra_headers=extra_headers if extra_headers else None,
|
||||
n=1,
|
||||
**token_params,
|
||||
)
|
||||
|
||||
@override
|
||||
def chat(
|
||||
self,
|
||||
messages: list[LLMMessage],
|
||||
model_config: ModelConfig,
|
||||
tools: list[Tool] | None = None,
|
||||
reuse_history: bool = True,
|
||||
) -> LLMResponse:
|
||||
"""Send chat messages with optional tool support."""
|
||||
parsed_messages = self.parse_messages(messages)
|
||||
if reuse_history:
|
||||
self.message_history = self.message_history + parsed_messages
|
||||
else:
|
||||
self.message_history = parsed_messages
|
||||
|
||||
tool_schemas = None
|
||||
if tools:
|
||||
tool_schemas = [
|
||||
ChatCompletionToolParam(
|
||||
function=FunctionDefinition(
|
||||
name=tool.get_name(),
|
||||
description=tool.get_description(),
|
||||
parameters=tool.get_input_schema(),
|
||||
),
|
||||
type="function",
|
||||
)
|
||||
for tool in tools
|
||||
]
|
||||
|
||||
# Get provider-specific extra headers
|
||||
extra_headers = self.provider_config.get_extra_headers()
|
||||
|
||||
# Apply retry decorator to the API call
|
||||
retry_decorator = retry_with(
|
||||
func=self._create_response,
|
||||
provider_name=self.provider_config.get_service_name(),
|
||||
max_retries=model_config.max_retries,
|
||||
)
|
||||
response = retry_decorator(model_config, tool_schemas, extra_headers)
|
||||
|
||||
choice = response.choices[0]
|
||||
|
||||
tool_calls: list[ToolCall] | None = None
|
||||
if choice.message.tool_calls:
|
||||
tool_calls = []
|
||||
for tool_call in choice.message.tool_calls:
|
||||
tool_calls.append(
|
||||
ToolCall(
|
||||
name=tool_call.function.name,
|
||||
call_id=tool_call.id,
|
||||
arguments=(
|
||||
json.loads(tool_call.function.arguments)
|
||||
if tool_call.function.arguments
|
||||
else {}
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
llm_response = LLMResponse(
|
||||
content=choice.message.content or "",
|
||||
tool_calls=tool_calls,
|
||||
finish_reason=choice.finish_reason,
|
||||
model=response.model,
|
||||
usage=(
|
||||
LLMUsage(
|
||||
input_tokens=response.usage.prompt_tokens or 0,
|
||||
output_tokens=response.usage.completion_tokens or 0,
|
||||
)
|
||||
if response.usage
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
# Update message history
|
||||
if llm_response.tool_calls:
|
||||
self.message_history.append(
|
||||
ChatCompletionAssistantMessageParam(
|
||||
role="assistant",
|
||||
content=llm_response.content,
|
||||
tool_calls=[
|
||||
ChatCompletionMessageToolCallParam(
|
||||
id=tool_call.call_id,
|
||||
function=Function(
|
||||
name=tool_call.name,
|
||||
arguments=json.dumps(tool_call.arguments),
|
||||
),
|
||||
type="function",
|
||||
)
|
||||
for tool_call in llm_response.tool_calls
|
||||
],
|
||||
)
|
||||
)
|
||||
elif llm_response.content:
|
||||
self.message_history.append(
|
||||
ChatCompletionAssistantMessageParam(content=llm_response.content, role="assistant")
|
||||
)
|
||||
|
||||
if self.trajectory_recorder:
|
||||
self.trajectory_recorder.record_llm_interaction(
|
||||
messages=messages,
|
||||
response=llm_response,
|
||||
provider=self.provider_config.get_provider_name(),
|
||||
model=model_config.model,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
return llm_response
|
||||
|
||||
def parse_messages(self, messages: list[LLMMessage]) -> list[ChatCompletionMessageParam]:
|
||||
"""Parse LLM messages to OpenAI format."""
|
||||
openai_messages: list[ChatCompletionMessageParam] = []
|
||||
for msg in messages:
|
||||
match msg:
|
||||
case msg if msg.tool_call is not None:
|
||||
_msg_tool_call_handler(openai_messages, msg)
|
||||
case msg if msg.tool_result is not None:
|
||||
_msg_tool_result_handler(openai_messages, msg)
|
||||
case _:
|
||||
_msg_role_handler(openai_messages, msg)
|
||||
|
||||
return openai_messages
|
||||
|
||||
|
||||
def _msg_tool_call_handler(messages: list[ChatCompletionMessageParam], msg: LLMMessage) -> None:
|
||||
if msg.tool_call:
|
||||
messages.append(
|
||||
ChatCompletionFunctionMessageParam(
|
||||
content=json.dumps(
|
||||
{
|
||||
"name": msg.tool_call.name,
|
||||
"arguments": msg.tool_call.arguments,
|
||||
}
|
||||
),
|
||||
role="function",
|
||||
name=msg.tool_call.name,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _msg_tool_result_handler(messages: list[ChatCompletionMessageParam], msg: LLMMessage) -> None:
|
||||
if msg.tool_result:
|
||||
result: str = ""
|
||||
if msg.tool_result.result:
|
||||
result = result + msg.tool_result.result + "\n"
|
||||
if msg.tool_result.error:
|
||||
result += "Tool call failed with error:\n"
|
||||
result += msg.tool_result.error
|
||||
result = result.strip()
|
||||
messages.append(
|
||||
ChatCompletionToolMessageParam(
|
||||
content=result,
|
||||
role="tool",
|
||||
tool_call_id=msg.tool_result.call_id,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _msg_role_handler(messages: list[ChatCompletionMessageParam], msg: LLMMessage) -> None:
|
||||
if msg.role:
|
||||
match msg.role:
|
||||
case "system":
|
||||
if not msg.content:
|
||||
raise ValueError("System message content is required")
|
||||
messages.append(
|
||||
ChatCompletionSystemMessageParam(content=msg.content, role="system")
|
||||
)
|
||||
case "user":
|
||||
if not msg.content:
|
||||
raise ValueError("User message content is required")
|
||||
messages.append(ChatCompletionUserMessageParam(content=msg.content, role="user"))
|
||||
case "assistant":
|
||||
if not msg.content:
|
||||
raise ValueError("Assistant message content is required")
|
||||
messages.append(
|
||||
ChatCompletionAssistantMessageParam(content=msg.content, role="assistant")
|
||||
)
|
||||
case _:
|
||||
raise ValueError(f"Invalid message role: {msg.role}")
|
||||
@@ -0,0 +1,74 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""OpenRouter provider configuration."""
|
||||
|
||||
import os
|
||||
|
||||
import openai
|
||||
|
||||
from trae_agent.utils.config import ModelConfig
|
||||
from trae_agent.utils.llm_clients.openai_compatible_base import (
|
||||
OpenAICompatibleClient,
|
||||
ProviderConfig,
|
||||
)
|
||||
|
||||
|
||||
class OpenRouterProvider(ProviderConfig):
|
||||
"""OpenRouter provider configuration."""
|
||||
|
||||
def create_client(
|
||||
self, api_key: str, base_url: str | None, api_version: str | None
|
||||
) -> openai.OpenAI:
|
||||
"""Create OpenAI client with OpenRouter base URL."""
|
||||
return openai.OpenAI(api_key=api_key, base_url=base_url)
|
||||
|
||||
def get_service_name(self) -> str:
|
||||
"""Get the service name for retry logging."""
|
||||
return "OpenRouter"
|
||||
|
||||
def get_provider_name(self) -> str:
|
||||
"""Get the provider name for trajectory recording."""
|
||||
return "openrouter"
|
||||
|
||||
def get_extra_headers(self) -> dict[str, str]:
|
||||
"""Get OpenRouter-specific headers."""
|
||||
extra_headers: dict[str, str] = {}
|
||||
|
||||
openrouter_site_url = os.getenv("OPENROUTER_SITE_URL")
|
||||
if openrouter_site_url:
|
||||
extra_headers["HTTP-Referer"] = openrouter_site_url
|
||||
|
||||
openrouter_site_name = os.getenv("OPENROUTER_SITE_NAME")
|
||||
if openrouter_site_name:
|
||||
extra_headers["X-Title"] = openrouter_site_name
|
||||
|
||||
return extra_headers
|
||||
|
||||
def supports_tool_calling(self, model_name: str) -> bool:
|
||||
"""Check if the model supports tool calling."""
|
||||
# Most modern models on OpenRouter support tool calling
|
||||
# We'll be conservative and check for known capable models
|
||||
tool_capable_patterns = [
|
||||
"gpt-4",
|
||||
"gpt-3.5-turbo",
|
||||
"claude-3",
|
||||
"claude-2",
|
||||
"gemini",
|
||||
"mistral",
|
||||
"llama-3",
|
||||
"command-r",
|
||||
]
|
||||
return any(pattern in model_name.lower() for pattern in tool_capable_patterns)
|
||||
|
||||
|
||||
class OpenRouterClient(OpenAICompatibleClient):
|
||||
"""OpenRouter client wrapper that maintains compatibility while using the new architecture."""
|
||||
|
||||
def __init__(self, model_config: ModelConfig):
|
||||
if (
|
||||
model_config.model_provider.base_url is None
|
||||
or model_config.model_provider.base_url == ""
|
||||
):
|
||||
model_config.model_provider.base_url = "https://openrouter.ai/api/v1"
|
||||
super().__init__(model_config, OpenRouterProvider())
|
||||
@@ -0,0 +1,2 @@
|
||||
# Utils/models
|
||||
Refactor the list of models into a more robust and developer-friendly format.
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user