diff --git a/.gitattributes b/.gitattributes index 44948bef92cd..75290afcc5c7 100644 --- a/.gitattributes +++ b/.gitattributes @@ -91,5 +91,5 @@ python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/protos/*.pyi linguist- python/packages/autogen-ext/tests/protos/*.py linguist-generated python/packages/autogen-ext/tests/protos/*.pyi linguist-generated docs/** linguist-documentation -python/packages/autogen-core/docs/** linguist-documentation +python/docs/** linguist-documentation dotnet/website/** linguist-documentation diff --git a/.github/ISSUE_TEMPLATE/1-bug_report.yml b/.github/ISSUE_TEMPLATE/1-bug_report.yml index 7674be913db0..04055eb6b47c 100644 --- a/.github/ISSUE_TEMPLATE/1-bug_report.yml +++ b/.github/ISSUE_TEMPLATE/1-bug_report.yml @@ -90,6 +90,21 @@ body: multiple: false options: - "Python dev (main branch)" + - "Python 0.7.2" + - "Python 0.7.1" + - "Python 0.6.4" + - "Python 0.6.2" + - "Python 0.6.1" + - "Python 0.6.0" + - "Python 0.5.7" + - "Python 0.5.6" + - "Python 0.5.5" + - "Python 0.5.4" + - "Python 0.5.3" + - "Python 0.5.2" + - "Python 0.5.1" + - "Python 0.4.9" + - "Python 0.4.8" - "Python 0.4.7" - "Python 0.4.6" - "Python 0.4.5" diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 99714fe7ec96..d27fc2e421d5 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -18,7 +18,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" @@ -37,7 +36,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" @@ -66,7 +64,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" @@ -78,6 +75,24 @@ jobs: poe --directory ${{ matrix.package }} mypy working-directory: ./python + docs-mypy: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + - run: uv sync --locked --all-extras + working-directory: ./python + - name: Run task + run: | + source ${{ github.workspace }}/python/.venv/bin/activate + poe docs-mypy + working-directory: ./python + pyright: runs-on: ubuntu-latest strategy: @@ -95,7 +110,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" @@ -122,7 +136,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" @@ -156,7 +169,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" @@ -189,7 +201,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: @@ -251,15 +262,11 @@ jobs: docs: runs-on: ubuntu-latest - strategy: - matrix: - package: ["./packages/autogen-core"] steps: - uses: actions/checkout@v4 - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" @@ -268,20 +275,16 @@ jobs: - name: Run task run: | source ${{ github.workspace }}/python/.venv/bin/activate - poe --directory ${{ matrix.package }} docs-check + poe docs-check working-directory: ./python docs-example-check: runs-on: ubuntu-latest - strategy: - matrix: - package: ["./packages/autogen-core"] steps: - uses: actions/checkout@v4 - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" @@ -290,7 +293,7 @@ jobs: - name: Run task run: | source ${{ github.workspace }}/python/.venv/bin/activate - poe --directory ${{ matrix.package }} docs-check-examples + poe docs-check-examples working-directory: ./python samples-code-check: @@ -300,7 +303,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" @@ -319,7 +321,6 @@ jobs: - uses: astral-sh/setup-uv@v3 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" @@ -338,7 +339,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index 8a76bde68948..19a6ac79ac4b 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -76,7 +76,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - run: uv sync --locked --all-extras working-directory: ./python - name: Prepare python venv diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index e936adf3750c..de0649d92761 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -31,19 +31,174 @@ jobs: matrix: version: [ - # For main use the workflow target - { ref: "${{github.ref}}", dest-dir: dev, uv-version: "0.5.13", sphinx-release-override: "dev" }, - { ref: "python-v0.4.9-website", dest-dir: stable, uv-version: "0.5.13", sphinx-release-override: "stable" }, - { ref: "v0.4.0.post1", dest-dir: "0.4.0", uv-version: "0.5.13", sphinx-release-override: "" }, - { ref: "v0.4.1", dest-dir: "0.4.1", uv-version: "0.5.13", sphinx-release-override: "" }, - { ref: "v0.4.2", dest-dir: "0.4.2", uv-version: "0.5.13", sphinx-release-override: "" }, - { ref: "v0.4.3", dest-dir: "0.4.3", uv-version: "0.5.13", sphinx-release-override: "" }, - { ref: "v0.4.4", dest-dir: "0.4.4", uv-version: "0.5.13", sphinx-release-override: "" }, - { ref: "python-v0.4.5", dest-dir: "0.4.5", uv-version: "0.5.13", sphinx-release-override: "" }, - { ref: "python-v0.4.6", dest-dir: "0.4.6", uv-version: "0.5.13", sphinx-release-override: "" }, - { ref: "python-v0.4.7", dest-dir: "0.4.7", uv-version: "0.5.13", sphinx-release-override: "" }, - { ref: "python-v0.4.8", dest-dir: "0.4.8", uv-version: "0.5.13", sphinx-release-override: "" }, - { ref: "python-v0.4.9-website", dest-dir: "0.4.9", uv-version: "0.5.13", sphinx-release-override: "" }, + { + ref: "${{github.ref}}", + dest-dir: dev, + uv-version: "0.7.13", + sphinx-release-override: "dev", + poe-dir: ".", + }, + { + ref: "python-v0.7.2", + dest-dir: stable, + uv-version: "0.7.13", + sphinx-release-override: "stable", + poe-dir: ".", + }, + { + ref: "v0.4.0.post1", + dest-dir: "0.4.0", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "v0.4.1", + dest-dir: "0.4.1", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "v0.4.2", + dest-dir: "0.4.2", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "v0.4.3", + dest-dir: "0.4.3", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "v0.4.4", + dest-dir: "0.4.4", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.4.5", + dest-dir: "0.4.5", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.4.6", + dest-dir: "0.4.6", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.4.7", + dest-dir: "0.4.7", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.4.8", + dest-dir: "0.4.8", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.4.9-website", + dest-dir: "0.4.9", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.5.1", + dest-dir: "0.5.1", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.5.2", + dest-dir: "0.5.2", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.5.3", + dest-dir: "0.5.3", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.5.4", + dest-dir: "0.5.4", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.5.5", + dest-dir: "0.5.5", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.5.6", + dest-dir: "0.5.6", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.5.7", + dest-dir: "0.5.7", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.6.1", + dest-dir: "0.6.1", + uv-version: "0.5.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.6.2", + dest-dir: "0.6.2", + uv-version: "0.7.13", + sphinx-release-override: "", + poe-dir: "./packages/autogen-core", + }, + { + ref: "python-v0.6.4", + dest-dir: "0.6.4", + uv-version: "0.7.13", + sphinx-release-override: "", + poe-dir: ".", + }, + { + ref: "python-v0.7.1.post1", + dest-dir: "0.7.1", + uv-version: "0.7.13", + sphinx-release-override: "", + poe-dir: ".", + }, + { + ref: "python-v0.7.2", + dest-dir: "0.7.2", + uv-version: "0.7.13", + sphinx-release-override: "", + poe-dir: ".", + }, ] steps: - name: Checkout @@ -59,11 +214,14 @@ jobs: with: python-version: "3.11" - run: | - uv sync --locked --all-extras + uv venv --python=3.11 source .venv/bin/activate - poe --directory ./packages/autogen-core docs-build + # Only pin version to ensure compatibility + uv pip install -U 'setuptools-scm>=8.1' + uv sync --locked --all-extras + poe --directory ${{ matrix.version.poe-dir }} docs-build mkdir -p docs-staging/${{ matrix.version.dest-dir }}/ - mv ./packages/autogen-core/docs/build/* docs-staging/${{ matrix.version.dest-dir }}/ + mv ${{ matrix.version.poe-dir }}/docs/build/* docs-staging/${{ matrix.version.dest-dir }}/ working-directory: ./python env: PY_DOCS_DIR: ${{ matrix.version.dest-dir }}/ @@ -87,7 +245,7 @@ jobs: - name: generate redirects run: | mkdir -p python/docs-staging/ - python python/packages/autogen-core/docs/redirects/redirects.py python/docs-staging + python python/docs/redirects/redirects.py python/docs-staging - uses: actions/upload-artifact@v4 with: path: "./python/docs-staging" @@ -104,7 +262,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" @@ -135,10 +292,11 @@ jobs: - name: setup python uses: actions/setup-python@v5 with: - python-version: "3.8" + python-version: "3.9" - name: pydoc-markdown install run: | python -m pip install --upgrade pip + pip install docspec==2.2.1 docspec-python==2.2.1 pip install pydoc-markdown pyyaml termcolor # Pin databind packages as version 4.5.0 is not compatible with pydoc-markdown. pip install databind.core==4.4.2 databind.json==4.4.2 @@ -232,7 +390,8 @@ jobs: name: github-pages url: ${{ steps.deployment.outputs.page_url }} runs-on: ubuntu-latest - needs: [build-02, build-04, build-04-dotnet, gen-redirects, gen-component-schema] + needs: + [build-02, build-04, build-04-dotnet, gen-redirects, gen-component-schema] if: ${{ needs.build-02.result == 'success' && needs.build-04.result == 'success' && needs.gen-redirects.result == 'success' && github.ref == 'refs/heads/main' }} steps: - uses: actions/download-artifact@v4 diff --git a/.github/workflows/dotnet-build.yml b/.github/workflows/dotnet-build.yml index 3d83bdf39d19..a5145dbfa80e 100644 --- a/.github/workflows/dotnet-build.yml +++ b/.github/workflows/dotnet-build.yml @@ -68,7 +68,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" @@ -154,7 +153,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" @@ -247,7 +245,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 4dfdfa6f1e3a..39a29f2efdd5 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -27,7 +27,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - uses: actions/setup-python@v5 with: python-version: "3.11" diff --git a/.github/workflows/pytest-mem0.yml b/.github/workflows/pytest-mem0.yml new file mode 100644 index 000000000000..75f177e1a0a9 --- /dev/null +++ b/.github/workflows/pytest-mem0.yml @@ -0,0 +1,94 @@ +name: Mem0 Memory Tests + +on: + # Run on pushes to any branch + push: + # Also run on pull requests to main + pull_request: + branches: + - main + +jobs: + test: + runs-on: ubuntu-latest + + services: + neo4j: + image: neo4j:5.26.6 + ports: + - 7474:7474 # HTTP + - 7687:7687 # BOLT + env: + NEO4J_AUTH: neo4j/password + NEO4J_dbms_security_procedures_unrestricted: apoc.* + # Add this to ensure Neo4j is ready for connections quickly + NEO4J_dbms_memory_pagecache_size: 100M + NEO4J_dbms_memory_heap_initial__size: 100M + NEO4J_dbms_memory_heap_max__size: 500M + # Try a different health check approach + options: >- + --health-cmd "wget -O /dev/null -q http://localhost:7474 || exit 1" + --health-interval 5s + --health-timeout 15s + --health-retries 10 + --health-start-period 30s + + steps: + - name: Check out repository + uses: actions/checkout@v3 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.11' + + - name: Wait for Neo4j + run: | + # Give Neo4j some extra time to start up + sleep 10 + # Try to connect to Neo4j + timeout 30s bash -c 'until curl -s http://localhost:7474 > /dev/null; do sleep 1; done' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + + # Install core packages first (in the right order) + cd python/packages/autogen-core + pip install -e . + + cd ../autogen-agentchat + pip install -e . + + # Now install autogen-ext with its dependencies + cd ../autogen-ext + pip install -e ".[dev,mem0,mem0-local]" + + # Install test dependencies + pip install pytest pytest-asyncio pytest-cov + pip install python-dotenv + + # Install dependencies for complex configuration tests + pip install "openai>=1.0.0" + pip install deepseek-ai + + # Update test config to match the simplified Neo4j setup + - name: Update Neo4j password in tests + run: | + echo "NEO4J_PASSWORD=password" >> $GITHUB_ENV + + - name: Run tests with coverage + # env: + # MEM0_API_KEY: ${{ secrets.MEM0_API_KEY }} + # SF_API_KEY: ${{ secrets.SF_API_KEY }} + # DEEPSEEK_API_KEY: ${{ secrets.DEEPSEEK_API_KEY }} + run: | + cd python/packages/autogen-ext + pytest --cov=autogen_ext.memory.mem0 tests/memory/test_mem0.py -v --cov-report=xml + + - name: Upload coverage to Codecov + uses: codecov/codecov-action@v3 + with: + file: ./python/packages/autogen-ext/coverage.xml + name: codecov-mem0 + fail_ci_if_error: false diff --git a/.github/workflows/pytest-redis-memory.yml b/.github/workflows/pytest-redis-memory.yml new file mode 100644 index 000000000000..bce5ba52ef30 --- /dev/null +++ b/.github/workflows/pytest-redis-memory.yml @@ -0,0 +1,67 @@ +name: Redis Memory Tests + +on: + push: + pull_request: + branches: + - main + +jobs: + test: + runs-on: ubuntu-latest + + services: + redis: + image: redis:latest + ports: + - 6379:6379 + env: + REDIS_URL: redis://localhost:6379 + + steps: + - name: Check out repository + uses: actions/checkout@v3 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.11' + + - name: Wait for Redis + run: | + # give Redis time to start + sleep 5 + # Wait for Redis to respond to curl (expecting empty reply, code 52) + timeout 5s bash -c 'until curl -s localhost:6379 || [ $? -eq 52 ]; do sleep 1; done' + - name: Install dependencies + run: | + python -m pip install --upgrade pip + + # install core packages + cd python/packages/autogen-core + pip install -e . + + cd ../autogen-agentchat + pip install -e . + + # install autogen-ext with its dependencies + cd ../autogen-ext + pip install -e ".[dev,redisvl]" + + # install test dependencies + pip install pytest pytest-asyncio pytest-cov + + # install additional dependencies for redis memory tests + pip install sentence-transformers + + - name: Run tests with coverage + run: | + cd python/packages/autogen-ext + pytest --cov=autogen_ext.memory.redis tests/memory/test_redis_memory.py -v --cov-report=xml + + - name: Upload coverage to Codecov + uses: codecov/codecov-action@v3 + with: + file: ./python/packages/autogen-ext/coverage.xml + name: codecov-redis-memory + fail_ci_if_error: false diff --git a/.github/workflows/single-python-package.yml b/.github/workflows/single-python-package.yml index 7a4d8bcafb0a..4d46ed5f8175 100644 --- a/.github/workflows/single-python-package.yml +++ b/.github/workflows/single-python-package.yml @@ -4,7 +4,7 @@ on: workflow_dispatch: inputs: package: - description: 'Select the package to deploy' + description: "Select the package to deploy" required: true type: choice options: @@ -14,8 +14,9 @@ on: - agbench - autogen-studio - magentic-one-cli + - pyautogen ref: - description: 'Tag to deploy' + description: "Tag to deploy" required: true jobs: @@ -35,7 +36,6 @@ jobs: - uses: astral-sh/setup-uv@v5 with: enable-cache: true - version: "0.5.18" - run: uv build --package ${{ github.event.inputs.package }} --out-dir dist/ working-directory: python - name: Publish package to PyPI diff --git a/.gitignore b/.gitignore index 53c18fbd1b6c..a008a99c4e96 100644 --- a/.gitignore +++ b/.gitignore @@ -202,3 +202,4 @@ registry.json # files created by the gitty agent in python/samples/gitty .gitty/ +.aider* diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c2e781239c48..ba1168b608d6 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -34,7 +34,7 @@ For common tasks that are helpful during development and run in CI, see [here](. ## Roadmap -We use GitHub issues and milestones to track our roadmap. You can view the upcoming milestones [here]([Roadmap Issues](https://aka.ms/autogen-roadmap). +We use GitHub issues and milestones to track our roadmap. You can view the upcoming milestones [here]([Roadmap Issues](https://aka.ms/autogen-roadmap)). ## Versioning @@ -48,11 +48,11 @@ We will update verion numbers according to the following rules: ## Release process 1. Create a PR that updates the version numbers across the codebase ([example](https://github.com/microsoft/autogen/pull/4359)) - 2. The docs CI will fail for the PR, but this is expected and will be resolved in the next step -2. After merging the PR, create and push a tag that corresponds to the new verion. For example, for `0.4.0.dev13`: +2. The docs CI will fail for the PR, but this is expected and will be resolved in the next step +3. After merging the PR, create and push a tag that corresponds to the new verion. For example, for `0.4.0.dev13`: - `git tag v0.4.0.dev13 && git push origin v0.4.0.dev13` -3. Restart the docs CI by finding the failed [job corresponding to the `push` event](https://github.com/microsoft/autogen/actions/workflows/docs.yml) and restarting all jobs -4. Run [this](https://github.com/microsoft/autogen/actions/workflows/single-python-package.yml) workflow for each of the packages that need to be released and get an approval for the release for it to run +4. Restart the docs CI by finding the failed [job corresponding to the `push` event](https://github.com/microsoft/autogen/actions/workflows/docs.yml) and restarting all jobs +5. Run [this](https://github.com/microsoft/autogen/actions/workflows/single-python-package.yml) workflow for each of the packages that need to be released and get an approval for the release for it to run ## Triage process diff --git a/README.md b/README.md index 73f4cd309469..2824058aa200 100644 --- a/README.md +++ b/README.md @@ -11,10 +11,6 @@ -
- Important: This is the official project. We are not affiliated with any fork or startup. See our statement. -
- # AutoGen **AutoGen** is a framework for creating multi-agent AI applications that can act autonomously or work alongside humans. @@ -55,43 +51,51 @@ async def main() -> None: asyncio.run(main()) ``` -### Web Browsing Agent Team +### MCP Server -Create a group chat team with a web surfer agent and a user proxy agent -for web browsing tasks. You need to install [playwright](https://playwright.dev/python/docs/library). +Create a web browsing assistant agent that uses the Playwright MCP server. ```python -# pip install -U autogen-agentchat autogen-ext[openai,web-surfer] -# playwright install +# First run `npm install -g @playwright/mcp@latest` to install the MCP server. import asyncio -from autogen_agentchat.agents import UserProxyAgent -from autogen_agentchat.conditions import TextMentionTermination +from autogen_agentchat.agents import AssistantAgent from autogen_agentchat.teams import RoundRobinGroupChat +from autogen_agentchat.conditions import TextMessageTermination from autogen_agentchat.ui import Console from autogen_ext.models.openai import OpenAIChatCompletionClient -from autogen_ext.agents.web_surfer import MultimodalWebSurfer +from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams + async def main() -> None: - model_client = OpenAIChatCompletionClient(model="gpt-4o") - # The web surfer will open a Chromium browser window to perform web browsing tasks. - web_surfer = MultimodalWebSurfer("web_surfer", model_client, headless=False, animate_actions=True) - # The user proxy agent is used to get user input after each step of the web surfer. - # NOTE: you can skip input by pressing Enter. - user_proxy = UserProxyAgent("user_proxy") - # The termination condition is set to end the conversation when the user types 'exit'. - termination = TextMentionTermination("exit", sources=["user_proxy"]) - # Web surfer and user proxy take turns in a round-robin fashion. - team = RoundRobinGroupChat([web_surfer, user_proxy], termination_condition=termination) - try: - # Start the team and wait for it to terminate. - await Console(team.run_stream(task="Find information about AutoGen and write a short summary.")) - finally: - await web_surfer.close() - await model_client.close() + model_client = OpenAIChatCompletionClient(model="gpt-4.1") + server_params = StdioServerParams( + command="npx", + args=[ + "@playwright/mcp@latest", + "--headless", + ], + ) + async with McpWorkbench(server_params) as mcp: + agent = AssistantAgent( + "web_browsing_assistant", + model_client=model_client, + workbench=mcp, # For multiple MCP servers, put them in a list. + model_client_stream=True, + max_tool_iterations=10, + ) + team = RoundRobinGroupChat( + [agent], + termination_condition=TextMessageTermination(source="web_browsing_assistant"), + ) + await Console(team.run_stream(task="Find out how many contributors for the microsoft/autogen repository")) + asyncio.run(main()) ``` +> **Warning**: Only connect to trusted MCP servers as they may execute commands +> in your local environment or expose sensitive information. + ### AutoGen Studio Use AutoGen Studio to prototype and run multi-agent workflows without writing code. @@ -132,17 +136,16 @@ With AutoGen you get to join and contribute to a thriving ecosystem. We host wee
-| | [![Python](https://img.shields.io/badge/AutoGen-Python-blue?logo=python&logoColor=white)](./python) | [![.NET](https://img.shields.io/badge/AutoGen-.NET-green?logo=.net&logoColor=white)](./dotnet) | [![Studio](https://img.shields.io/badge/AutoGen-Studio-purple?logo=visual-studio&logoColor=white)](./python/packages/autogen-studio) | -| ------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------- | -| Installation | [![Installation](https://img.shields.io/badge/Install-blue)](https://microsoft.github.io/autogen/stable/user-guide/agentchat-user-guide/installation.html) | [![Install](https://img.shields.io/badge/Install-green)](https://microsoft.github.io/autogen/dotnet/dev/core/installation.html) | [![Install](https://img.shields.io/badge/Install-purple)](https://microsoft.github.io/autogen/stable/user-guide/autogenstudio-user-guide/installation.html) | -| Quickstart | [![Quickstart](https://img.shields.io/badge/Quickstart-blue)](https://microsoft.github.io/autogen/stable/user-guide/agentchat-user-guide/quickstart.html#) | [![Quickstart](https://img.shields.io/badge/Quickstart-green)](https://microsoft.github.io/autogen/dotnet/dev/core/index.html) | [![Usage](https://img.shields.io/badge/Quickstart-purple)](https://microsoft.github.io/autogen/stable/user-guide/autogenstudio-user-guide/usage.html#) | -| Tutorial | [![Tutorial](https://img.shields.io/badge/Tutorial-blue)](https://microsoft.github.io/autogen/stable/user-guide/agentchat-user-guide/tutorial/index.html) | [![Tutorial](https://img.shields.io/badge/Tutorial-green)](https://microsoft.github.io/autogen/dotnet/dev/core/tutorial.html) | [![Usage](https://img.shields.io/badge/Tutorial-purple)](https://microsoft.github.io/autogen/stable/user-guide/autogenstudio-user-guide/usage.html#) | -| API Reference | [![API](https://img.shields.io/badge/Docs-blue)](https://microsoft.github.io/autogen/stable/reference/index.html#) | [![API](https://img.shields.io/badge/Docs-green)](https://microsoft.github.io/autogen/dotnet/dev/api/Microsoft.AutoGen.Contracts.html) | [![API](https://img.shields.io/badge/Docs-purple)](https://microsoft.github.io/autogen/stable/user-guide/autogenstudio-user-guide/usage.html) | -| Packages | [![PyPi autogen-core](https://img.shields.io/badge/PyPi-autogen--core-blue?logo=pypi)](https://pypi.org/project/autogen-core/)
[![PyPi autogen-agentchat](https://img.shields.io/badge/PyPi-autogen--agentchat-blue?logo=pypi)](https://pypi.org/project/autogen-agentchat/)
[![PyPi autogen-ext](https://img.shields.io/badge/PyPi-autogen--ext-blue?logo=pypi)](https://pypi.org/project/autogen-ext/) | [![NuGet Contracts](https://img.shields.io/badge/NuGet-Contracts-green?logo=nuget)](https://www.nuget.org/packages/Microsoft.AutoGen.Contracts/)
[![NuGet Core](https://img.shields.io/badge/NuGet-Core-green?logo=nuget)](https://www.nuget.org/packages/Microsoft.AutoGen.Core/)
[![NuGet Core.Grpc](https://img.shields.io/badge/NuGet-Core.Grpc-green?logo=nuget)](https://www.nuget.org/packages/Microsoft.AutoGen.Core.Grpc/)
[![NuGet RuntimeGateway.Grpc](https://img.shields.io/badge/NuGet-RuntimeGateway.Grpc-green?logo=nuget)](https://www.nuget.org/packages/Microsoft.AutoGen.RuntimeGateway.Grpc/) | [![PyPi autogenstudio](https://img.shields.io/badge/PyPi-autogenstudio-purple?logo=pypi)](https://pypi.org/project/autogenstudio/) | +| | [![Python](https://img.shields.io/badge/AutoGen-Python-blue?logo=python&logoColor=white)](./python) | [![.NET](https://img.shields.io/badge/AutoGen-.NET-green?logo=.net&logoColor=white)](./dotnet) | [![Studio](https://img.shields.io/badge/AutoGen-Studio-purple?logo=visual-studio&logoColor=white)](./python/packages/autogen-studio) | +| ------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Installation | [![Installation](https://img.shields.io/badge/Install-blue)](https://microsoft.github.io/autogen/stable/user-guide/agentchat-user-guide/installation.html) | [![Install](https://img.shields.io/badge/Install-green)](https://microsoft.github.io/autogen/dotnet/dev/core/installation.html) | [![Install](https://img.shields.io/badge/Install-purple)](https://microsoft.github.io/autogen/stable/user-guide/autogenstudio-user-guide/installation.html) | +| Quickstart | [![Quickstart](https://img.shields.io/badge/Quickstart-blue)](https://microsoft.github.io/autogen/stable/user-guide/agentchat-user-guide/quickstart.html#) | [![Quickstart](https://img.shields.io/badge/Quickstart-green)](https://microsoft.github.io/autogen/dotnet/dev/core/index.html) | [![Usage](https://img.shields.io/badge/Quickstart-purple)](https://microsoft.github.io/autogen/stable/user-guide/autogenstudio-user-guide/usage.html#) | +| Tutorial | [![Tutorial](https://img.shields.io/badge/Tutorial-blue)](https://microsoft.github.io/autogen/stable/user-guide/agentchat-user-guide/tutorial/index.html) | [![Tutorial](https://img.shields.io/badge/Tutorial-green)](https://microsoft.github.io/autogen/dotnet/dev/core/tutorial.html) | [![Usage](https://img.shields.io/badge/Tutorial-purple)](https://microsoft.github.io/autogen/stable/user-guide/autogenstudio-user-guide/usage.html#) | +| API Reference | [![API](https://img.shields.io/badge/Docs-blue)](https://microsoft.github.io/autogen/stable/reference/index.html#) | [![API](https://img.shields.io/badge/Docs-green)](https://microsoft.github.io/autogen/dotnet/dev/api/Microsoft.AutoGen.Contracts.html) | [![API](https://img.shields.io/badge/Docs-purple)](https://microsoft.github.io/autogen/stable/user-guide/autogenstudio-user-guide/usage.html) | +| Packages | [![PyPi autogen-core](https://img.shields.io/badge/PyPi-autogen--core-blue?logo=pypi)](https://pypi.org/project/autogen-core/)
[![PyPi autogen-agentchat](https://img.shields.io/badge/PyPi-autogen--agentchat-blue?logo=pypi)](https://pypi.org/project/autogen-agentchat/)
[![PyPi autogen-ext](https://img.shields.io/badge/PyPi-autogen--ext-blue?logo=pypi)](https://pypi.org/project/autogen-ext/) | [![NuGet Contracts](https://img.shields.io/badge/NuGet-Contracts-green?logo=nuget)](https://www.nuget.org/packages/Microsoft.AutoGen.Contracts/)
[![NuGet Core](https://img.shields.io/badge/NuGet-Core-green?logo=nuget)](https://www.nuget.org/packages/Microsoft.AutoGen.Core/)
[![NuGet Core.Grpc](https://img.shields.io/badge/NuGet-Core.Grpc-green?logo=nuget)](https://www.nuget.org/packages/Microsoft.AutoGen.Core.Grpc/)
[![NuGet RuntimeGateway.Grpc](https://img.shields.io/badge/NuGet-RuntimeGateway.Grpc-green?logo=nuget)](https://www.nuget.org/packages/Microsoft.AutoGen.RuntimeGateway.Grpc/) | [![PyPi autogenstudio](https://img.shields.io/badge/PyPi-autogenstudio-purple?logo=pypi)](https://pypi.org/project/autogenstudio/) |
- Interested in contributing? See [CONTRIBUTING.md](./CONTRIBUTING.md) for guidelines on how to get started. We welcome contributions of all kinds, including bug fixes, new features, and documentation improvements. Join our community and help us make AutoGen better! Have questions? Check out our [Frequently Asked Questions (FAQ)](./FAQ.md) for answers to common queries. If you don't find what you're looking for, feel free to ask in our [GitHub Discussions](https://github.com/microsoft/autogen/discussions) or join our [Discord server](https://aka.ms/autogen-discord) for real-time support. You can also read our [blog](https://devblogs.microsoft.com/autogen/) for updates. diff --git a/docs/design/01 - Programming Model.md b/docs/design/01 - Programming Model.md index 6bfa9f9766ad..29705194456a 100644 --- a/docs/design/01 - Programming Model.md +++ b/docs/design/01 - Programming Model.md @@ -11,7 +11,7 @@ Each event in the system is defined using the [CloudEvents Specification](https: 1. *id* - A unique id (eg. a UUID). 2. *source* - A URI or URN indicating the event's origin. 3. *type* - The namespace of the event - prefixed with a reverse-DNS name. - - The prefixed domain dictates the organization which defines the semantics of this event type: e.g `com.github.pull_request.opened` or `com.example.object.deleted.v2`), and optionally fields describing the data schema/content-type or extensions. + - The prefixed domain dictates the organization which defines the semantics of this event type: e.g (`com.github.pull_request.opened` or `com.example.object.deleted.v2`), and optionally fields describing the data schema/content-type or extensions. ## Event Handlers diff --git a/docs/design/02 - Topics.md b/docs/design/02 - Topics.md index d7c93417cfc7..008e1aa9bfde 100644 --- a/docs/design/02 - Topics.md +++ b/docs/design/02 - Topics.md @@ -62,7 +62,7 @@ For this subscription source should map directly to agent key. This subscription will therefore receive all events for the following well known topics: -- `{AgentType}:` - General purpose direct messages. These should be routed to the approriate message handler. -- `{AgentType}:rpc_request={RequesterAgentType}` - RPC request messages. These should be routed to the approriate RPC handler, and RequesterAgentType used to publish the response +- `{AgentType}:` - General purpose direct messages. These should be routed to the appropriate message handler. +- `{AgentType}:rpc_request={RequesterAgentType}` - RPC request messages. These should be routed to the appropriate RPC handler, and RequesterAgentType used to publish the response - `{AgentType}:rpc_response={RequestId}` - RPC response messages. These should be routed back to the response future of the caller. - `{AgentType}:error={RequestId}` - Error message that corresponds to the given request. diff --git a/docs/switcher.json b/docs/switcher.json index 7a0e6d343a52..ec9bb130bd89 100644 --- a/docs/switcher.json +++ b/docs/switcher.json @@ -5,11 +5,71 @@ "url": "/autogen/dev/" }, { - "name": "0.4.9 (stable)", + "name": "0.7.2 (stable)", "version": "stable", "url": "/autogen/stable/", "preferred": true }, + { + "name": "0.7.1", + "version": "0.7.1", + "url": "/autogen/0.7.1/" + }, + { + "name": "0.6.4", + "version": "0.6.4", + "url": "/autogen/0.6.4/" + }, + { + "name": "0.6.2", + "version": "0.6.2", + "url": "/autogen/0.6.2/" + }, + { + "name": "0.6.1", + "version": "0.6.1", + "url": "/autogen/0.6.1/" + }, + { + "name": "0.5.7", + "version": "0.5.7", + "url": "/autogen/0.5.7/" + }, + { + "name": "0.5.6", + "version": "0.5.6", + "url": "/autogen/0.5.6/" + }, + { + "name": "0.5.5", + "version": "0.5.5", + "url": "/autogen/0.5.5/" + }, + { + "name": "0.5.4", + "version": "0.5.4", + "url": "/autogen/0.5.4/" + }, + { + "name": "0.5.3", + "version": "0.5.3", + "url": "/autogen/0.5.3/" + }, + { + "name": "0.5.2", + "version": "0.5.2", + "url": "/autogen/0.5.2/" + }, + { + "name": "0.5.1", + "version": "0.5.1", + "url": "/autogen/0.5.1/" + }, + { + "name": "0.4.9", + "version": "0.4.9", + "url": "/autogen/0.4.9/" + }, { "name": "0.4.8", "version": "0.4.8", @@ -60,4 +120,4 @@ "version": "0.2", "url": "/autogen/0.2/" } -] +] \ No newline at end of file diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index aca7a136ddf8..5240e16adab6 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -2,10 +2,13 @@ true 1.22.0 - 1.22.0-alpha - 9.0.0-preview.9.24525.1 + 1.45.0 + $(MicrosoftSemanticKernelStableVersion)-preview + $(MicrosoftSemanticKernelStableVersion)-alpha + 9.5.0 + 9.5.0-preview.1.25265.7 9.0.0 - 9.0.0 + 9.0.3 9.0.0 9.0.0 9.0.1 @@ -18,7 +21,7 @@ - + @@ -27,9 +30,10 @@ - + + @@ -61,9 +65,9 @@ - - - + + + @@ -100,18 +104,18 @@ - - - - - - + + + + + + - + @@ -127,9 +131,9 @@ - + - \ No newline at end of file + diff --git a/dotnet/README.md b/dotnet/README.md index 7647d227aa96..99b7ac3bfed1 100644 --- a/dotnet/README.md +++ b/dotnet/README.md @@ -40,7 +40,7 @@ var assistantAgent = new AssistantAgent( // set human input mode to ALWAYS so that user always provide input var userProxyAgent = new UserProxyAgent( name: "user", - humanInputMode: ConversableAgent.HumanInputMode.ALWAYS) + humanInputMode: HumanInputMode.ALWAYS) .RegisterPrintMessage(); // start the conversation diff --git a/dotnet/eng/MetaInfo.props b/dotnet/eng/MetaInfo.props index 7dd72c35c27a..169901b8b71c 100644 --- a/dotnet/eng/MetaInfo.props +++ b/dotnet/eng/MetaInfo.props @@ -2,7 +2,7 @@ 0.4.0 - 0.2.2 + 0.2.3 Microsoft https://microsoft.github.io/autogen-for-net/ https://github.com/microsoft/autogen diff --git a/dotnet/src/AutoGen.Core/Function/FunctionAttribute.cs b/dotnet/src/AutoGen.Core/Function/FunctionAttribute.cs index 9367f5c6f297..644c899c153f 100644 --- a/dotnet/src/AutoGen.Core/Function/FunctionAttribute.cs +++ b/dotnet/src/AutoGen.Core/Function/FunctionAttribute.cs @@ -4,6 +4,8 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Reflection; +using System.Text.Json; using System.Text.Json.Serialization; using Microsoft.Extensions.AI; @@ -69,36 +71,48 @@ public class FunctionContract /// public string? ReturnDescription { get; set; } - public static implicit operator FunctionContract(AIFunctionMetadata metadata) + public static implicit operator FunctionContract(AIFunction function) { - return new FunctionContract + var openapiScheme = function.JsonSchema; + var parameters = new List(); + string[] isRequiredProperties = []; + if (openapiScheme.TryGetProperty("required", out var requiredElement)) { - Namespace = metadata.AdditionalProperties.ContainsKey(NamespaceKey) ? metadata.AdditionalProperties[NamespaceKey] as string : null, - ClassName = metadata.AdditionalProperties.ContainsKey(ClassNameKey) ? metadata.AdditionalProperties[ClassNameKey] as string : null, - Name = metadata.Name, - Description = metadata.Description, - Parameters = metadata.Parameters?.Select(p => (FunctionParameterContract)p).ToList(), - ReturnType = metadata.ReturnParameter.ParameterType, - ReturnDescription = metadata.ReturnParameter.Description, - }; - } + isRequiredProperties = requiredElement.Deserialize() ?? []; + } - public static implicit operator AIFunctionMetadata(FunctionContract contract) - { - return new AIFunctionMetadata(contract.Name) + var parameterList = function.UnderlyingMethod?.GetParameters() ?? Array.Empty(); + + if (openapiScheme.TryGetProperty("properties", out var propertiesElement)) { - Description = contract.Description, - ReturnParameter = new AIFunctionReturnParameterMetadata() + var properties = propertiesElement.Deserialize>() ?? new Dictionary(); + foreach (var property in properties) { - Description = contract.ReturnDescription, - ParameterType = contract.ReturnType, - }, - AdditionalProperties = new Dictionary - { - [NamespaceKey] = contract.Namespace, - [ClassNameKey] = contract.ClassName, - }, - Parameters = [.. contract.Parameters?.Select(p => (AIFunctionParameterMetadata)p)!], + var parameterType = parameterList.FirstOrDefault(p => p.Name == property.Key)?.ParameterType; + var parameter = new FunctionParameterContract + { + Name = property.Key, + ParameterType = parameterType, // TODO: Need to get the type from the schema + IsRequired = isRequiredProperties.Contains(property.Key), + }; + if (property.Value.TryGetProperty("description", out var descriptionElement)) + { + parameter.Description = descriptionElement.GetString(); + } + if (property.Value.TryGetProperty("default", out var defaultValueElement)) + { + parameter.DefaultValue = defaultValueElement.Deserialize(); + } + parameters.Add(parameter); + } + } + return new FunctionContract + { + Namespace = function.AdditionalProperties.ContainsKey(NamespaceKey) ? function.AdditionalProperties[NamespaceKey] as string : null, + ClassName = function.AdditionalProperties.ContainsKey(ClassNameKey) ? function.AdditionalProperties[ClassNameKey] as string : null, + Name = function.Name, + Description = function.Description, + Parameters = parameters, }; } } @@ -132,29 +146,4 @@ public class FunctionParameterContract /// The default value of the parameter. /// public object? DefaultValue { get; set; } - - // convert to/from FunctionParameterMetadata - public static implicit operator FunctionParameterContract(AIFunctionParameterMetadata metadata) - { - return new FunctionParameterContract - { - Name = metadata.Name, - Description = metadata.Description, - ParameterType = metadata.ParameterType, - IsRequired = metadata.IsRequired, - DefaultValue = metadata.DefaultValue, - }; - } - - public static implicit operator AIFunctionParameterMetadata(FunctionParameterContract contract) - { - return new AIFunctionParameterMetadata(contract.Name!) - { - DefaultValue = contract.DefaultValue, - Description = contract.Description, - IsRequired = contract.IsRequired, - ParameterType = contract.ParameterType, - HasDefaultValue = contract.DefaultValue != null, - }; - } } diff --git a/dotnet/src/AutoGen.Core/Middleware/FunctionCallMiddleware.cs b/dotnet/src/AutoGen.Core/Middleware/FunctionCallMiddleware.cs index 266155316c81..60b9e703581d 100644 --- a/dotnet/src/AutoGen.Core/Middleware/FunctionCallMiddleware.cs +++ b/dotnet/src/AutoGen.Core/Middleware/FunctionCallMiddleware.cs @@ -53,9 +53,9 @@ public FunctionCallMiddleware( public FunctionCallMiddleware(IEnumerable functions, string? name = null) { this.Name = name ?? nameof(FunctionCallMiddleware); - this.functions = functions.Select(f => (FunctionContract)f.Metadata).ToArray(); + this.functions = functions.Select(f => (FunctionContract)f).ToArray(); - this.functionMap = functions.Select(f => (f.Metadata.Name, this.AIToolInvokeWrapper(f.InvokeAsync))).ToDictionary(f => f.Name, f => f.Item2); + this.functionMap = functions.Select(f => (f.Name, this.AIToolInvokeWrapper(f.InvokeAsync))).ToDictionary(f => f.Name, f => f.Item2); } public string? Name { get; } @@ -189,12 +189,12 @@ private async Task InvokeToolCallMessagesAfterInvokingAgentAsync(ToolC } } - private Func> AIToolInvokeWrapper(Func>?, CancellationToken, Task> lambda) + private Func> AIToolInvokeWrapper(Func> lambda) { return async (string args) => { var arguments = JsonSerializer.Deserialize>(args); - var result = await lambda(arguments, CancellationToken.None); + var result = await lambda(new(arguments), CancellationToken.None); return result switch { diff --git a/dotnet/src/AutoGen.SemanticKernel/SemanticKernelChatCompletionAgent.cs b/dotnet/src/AutoGen.SemanticKernel/SemanticKernelChatCompletionAgent.cs index 947807806976..2744bfdce00b 100644 --- a/dotnet/src/AutoGen.SemanticKernel/SemanticKernelChatCompletionAgent.cs +++ b/dotnet/src/AutoGen.SemanticKernel/SemanticKernelChatCompletionAgent.cs @@ -26,8 +26,9 @@ public SemanticKernelChatCompletionAgent(ChatCompletionAgent chatCompletionAgent public async Task GenerateReplyAsync(IEnumerable messages, GenerateReplyOptions? options = null, CancellationToken cancellationToken = default) { - ChatMessageContent[] reply = await _chatCompletionAgent - .InvokeAsync(BuildChatHistory(messages), cancellationToken: cancellationToken) + var agentThread = new ChatHistoryAgentThread(BuildChatHistory(messages)); + var reply = await _chatCompletionAgent + .InvokeAsync(agentThread, cancellationToken: cancellationToken) .ToArrayAsync(cancellationToken: cancellationToken); return reply.Length > 1 diff --git a/dotnet/src/Microsoft.AutoGen/AgentChat/Abstractions/Messages.cs b/dotnet/src/Microsoft.AutoGen/AgentChat/Abstractions/Messages.cs index 97e1d5b7628a..74edaf3e010c 100644 --- a/dotnet/src/Microsoft.AutoGen/AgentChat/Abstractions/Messages.cs +++ b/dotnet/src/Microsoft.AutoGen/AgentChat/Abstractions/Messages.cs @@ -124,7 +124,7 @@ public enum Type /// The to wrap. /// A instance wrapping the . /// - /// Thrown if the is not a or . + /// Thrown if the is not a or . /// public static MultiModalData CheckTypeAndCreate(AIContent item) { @@ -132,7 +132,7 @@ public static MultiModalData CheckTypeAndCreate(AIContent item) { return new MultiModalData(text); } - else if (item is ImageContent image) + else if (item is DataContent image) { return new MultiModalData(image); } @@ -163,10 +163,10 @@ public MultiModalData(TextContent textContent) } /// - /// Initializes a new instance of the with an . + /// Initializes a new instance of the with an . /// /// The image to wrap. - public MultiModalData(ImageContent image) + public MultiModalData(DataContent image) { ContentType = Type.Image; AIContent = image; @@ -254,12 +254,12 @@ public void AddRange(IEnumerable textItems) } /// - /// Adds a range of to the message. + /// Adds a range of to the message. /// /// The items to add. - public void AddRange(IEnumerable images) + public void AddRange(IEnumerable images) { - foreach (ImageContent image in images) + foreach (DataContent image in images) { this.Add(image); } @@ -287,7 +287,7 @@ public void Add(string text) /// Adds a to the message. /// /// The image to add. - public void Add(ImageContent image) + public void Add(DataContent image) { this.Content.Add(new(image)); } @@ -374,7 +374,7 @@ public void Insert(int index, TextContent text) } /// - public void Insert(int index, ImageContent image) + public void Insert(int index, DataContent image) { this.Content.Insert(index, new(image)); } @@ -610,7 +610,7 @@ public static Microsoft.Extensions.AI.ChatMessage Flatten(this Microsoft.Extensi { contentBuilder.AppendLine(textContent.Text); } - else if (content is ImageContent) + else if (content is DataContent) { contentBuilder.AppendLine("[Image]"); } diff --git a/dotnet/src/Microsoft.AutoGen/AgentChat/Abstractions/Tools.cs b/dotnet/src/Microsoft.AutoGen/AgentChat/Abstractions/Tools.cs index 1c0c276b7bbf..893c3392bc82 100644 --- a/dotnet/src/Microsoft.AutoGen/AgentChat/Abstractions/Tools.cs +++ b/dotnet/src/Microsoft.AutoGen/AgentChat/Abstractions/Tools.cs @@ -1,7 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Tools.cs -using System.ComponentModel; using System.Reflection; using Microsoft.Extensions.AI; @@ -10,37 +9,6 @@ namespace Microsoft.AutoGen.AgentChat.Abstractions; // TODO: This likely should live as a "Component" in an Agent-building ClassLib? // It seems like it could have applicability beyond AgentChat. -public static class ReflectionExtensions -{ - public static AIFunctionParameterMetadata ToAIFunctionMetadata(this ParameterInfo pi) - { - return new AIFunctionParameterMetadata(pi.Name!) - { - Description = pi.GetCustomAttribute()?.Description, - - ParameterType = pi.ParameterType, - - HasDefaultValue = pi.HasDefaultValue, - IsRequired = !pi.HasDefaultValue, - DefaultValue = pi.DefaultValue, - - // Schema = JSONSchema of type - }; - } - - public static AIFunctionReturnParameterMetadata ToAIFunctionReturnMetadata(this ParameterInfo rpi) - { - return new AIFunctionReturnParameterMetadata - { - Description = rpi.GetCustomAttribute()?.Description, - - ParameterType = rpi.ParameterType - - //Schema = JSONSchema of type - }; - } -} - public class ParameterSchema(string name, Type type, bool isRequired = false, object? defaultValue = default) { public string Name { get; } = name; @@ -54,15 +22,6 @@ public static implicit operator ParameterSchema(ParameterInfo parameterInfo) Type parameterType = parameterInfo.ParameterType; return ParameterSchema.Create(parameterType, parameterInfo.Name!, parameterInfo.HasDefaultValue, parameterInfo.DefaultValue); } - - public static implicit operator ParameterSchema(AIFunctionParameterMetadata parameterMetadata) - { - Type parameterType = parameterMetadata.ParameterType!; // TODO: Deal with missing ParameterTypes - return ParameterSchema.Create(parameterType, - parameterMetadata.Name, - parameterMetadata.IsRequired, - parameterMetadata.DefaultValue); - } } // TODO: Can this be obviated by AIFunctionParameter? @@ -86,7 +45,6 @@ public interface ITool public string Description { get; } public IEnumerable Parameters { get; } - public Type ReturnType { get; } // TODO: State serialization @@ -136,18 +94,15 @@ public class AIFunctionTool(AIFunction aiFunction) : ITool public AIFunction AIFunction { get; } = aiFunction; /// - public string Name => this.AIFunction.Metadata.Name; + public string Name => this.AIFunction.Name; /// - public string Description => this.AIFunction.Metadata.Description; + public string Description => this.AIFunction.Description; /// - public IEnumerable Parameters => from rawParameter in this.AIFunction.Metadata.Parameters + public IEnumerable Parameters => from rawParameter in this.AIFunction.UnderlyingMethod!.GetParameters() select (ParameterSchema)rawParameter; - /// - public Type ReturnType => this.AIFunction.Metadata.ReturnParameter.ParameterType!; // TODO: Deal with missing return types - /// public Task ExecuteAsync(IEnumerable parameters, CancellationToken cancellationToken = default) => this.ExecuteAsync(parameters, cancellationToken); @@ -164,23 +119,6 @@ public class CallableTool(string name, string description, Delegate callable) { internal static AIFunction CreateAIFunction(string name, string description, Delegate callable) { - MethodInfo methodInfo = callable.Method; - - IEnumerable parameters = - from parameterInfo in methodInfo.GetParameters() - select parameterInfo.ToAIFunctionMetadata(); - - AIFunctionReturnParameterMetadata returnParameter = methodInfo.ReturnParameter.ToAIFunctionReturnMetadata(); - - AIFunctionFactoryCreateOptions createOptions = new() - { - Name = name, - Description = description, - Parameters = parameters.ToList(), - ReturnParameter = returnParameter, - // SerializerOptions = TODO: How do we maintain consistency with Python? - }; - - return AIFunctionFactory.Create(callable, createOptions); + return AIFunctionFactory.Create(callable, name: name, description: description); } } diff --git a/dotnet/src/Microsoft.AutoGen/Agents/AIAgent/InferenceAgent.cs b/dotnet/src/Microsoft.AutoGen/Agents/AIAgent/InferenceAgent.cs index d3dc100012eb..8e753766b316 100644 --- a/dotnet/src/Microsoft.AutoGen/Agents/AIAgent/InferenceAgent.cs +++ b/dotnet/src/Microsoft.AutoGen/Agents/AIAgent/InferenceAgent.cs @@ -26,19 +26,19 @@ public abstract class InferenceAgent( { protected IChatClient ChatClient { get; } = client; private ILogger>? Logger => _logger as ILogger>; - private Task CompleteAsync( + private Task CompleteAsync( IList chatMessages, ChatOptions? options = null, CancellationToken cancellationToken = default) { - return ChatClient.CompleteAsync(chatMessages, options, cancellationToken); + return ChatClient.GetResponseAsync(chatMessages, options, cancellationToken); } - private IAsyncEnumerable CompleteStreamingAsync( + private IAsyncEnumerable CompleteStreamingAsync( IList chatMessages, ChatOptions? options = null, CancellationToken cancellationToken = default) { - return ChatClient.CompleteStreamingAsync(chatMessages, options, cancellationToken); + return ChatClient.GetStreamingResponseAsync(chatMessages, options, cancellationToken); } } diff --git a/dotnet/src/Microsoft.AutoGen/Extensions/MEAI/ServiceCollectionChatCompletionExtensions.cs b/dotnet/src/Microsoft.AutoGen/Extensions/MEAI/ServiceCollectionChatCompletionExtensions.cs index 114562993c29..28e5b414ef4d 100644 --- a/dotnet/src/Microsoft.AutoGen/Extensions/MEAI/ServiceCollectionChatCompletionExtensions.cs +++ b/dotnet/src/Microsoft.AutoGen/Extensions/MEAI/ServiceCollectionChatCompletionExtensions.cs @@ -41,12 +41,13 @@ public static IServiceCollection AddOllamaChatClient( Func? builder = null) { uri ??= new Uri("http://localhost:11434"); - return services.AddChatClient(pipeline => + services.AddChatClient(service => { - builder?.Invoke(pipeline); - var httpClient = pipeline.Services.GetService() ?? new(); - return pipeline.Use(new OllamaChatClient(uri, modelName, httpClient)); + var httpClient = service.GetService() ?? new(); + return new OllamaChatClient(uri, modelName, httpClient); }); + + return services; } public static IServiceCollection AddOpenAIChatClient( this IHostApplicationBuilder hostBuilder, @@ -81,16 +82,17 @@ public static IServiceCollection AddOpenAIChatClient( Uri? endpoint = null, Func? builder = null) { - return services + services .AddSingleton(_ => endpoint is null ? new OpenAIClient(apiKey) : new AzureOpenAIClient(endpoint, new ApiKeyCredential(apiKey))) - .AddChatClient(pipeline => + .AddChatClient(service => { - builder?.Invoke(pipeline); - var openAiClient = pipeline.Services.GetRequiredService(); - return pipeline.Use(openAiClient.AsChatClient(modelOrDeploymentName)); + var openAiClient = service.GetRequiredService(); + return openAiClient.GetChatClient(modelOrDeploymentName).AsIChatClient(); }); + + return services; } public static IServiceCollection AddAzureChatClient( this IHostApplicationBuilder hostBuilder, @@ -109,12 +111,10 @@ public static IServiceCollection AddAzureChatClient( } var endpoint = $"{serviceName}:Endpoint" ?? throw new InvalidOperationException($"No endpoint was specified for the Azure Inference Chat Client"); var endpointUri = string.IsNullOrEmpty(endpoint) ? null : new Uri(endpoint); - return hostBuilder.Services.AddChatClient(pipeline => - { - builder?.Invoke(pipeline); - var token = Environment.GetEnvironmentVariable("GH_TOKEN") ?? throw new InvalidOperationException("No model access token was found in the environment variable GH_TOKEN"); - return pipeline.Use(new ChatCompletionsClient( - endpointUri, new AzureKeyCredential(token)).AsChatClient(modelOrDeploymentName)); - }); + var token = Environment.GetEnvironmentVariable("AZURE_OPENAI_API_KEY") ?? throw new InvalidOperationException("No model access token was found in the environment variable AZURE_OPENAI_API_KEY"); + var chatClient = new ChatCompletionsClient(endpointUri, new AzureKeyCredential(token)).AsIChatClient(modelOrDeploymentName); + hostBuilder.Services.AddChatClient(chatClient); + + return hostBuilder.Services; } } diff --git a/dotnet/test/AutoGen.AzureAIInference.Tests/ChatCompletionClientAgentTests.cs b/dotnet/test/AutoGen.AzureAIInference.Tests/ChatCompletionClientAgentTests.cs index 5d2e0721ce64..5152c846165f 100644 --- a/dotnet/test/AutoGen.AzureAIInference.Tests/ChatCompletionClientAgentTests.cs +++ b/dotnet/test/AutoGen.AzureAIInference.Tests/ChatCompletionClientAgentTests.cs @@ -294,7 +294,7 @@ public async Task ItThrowExceptionWhenChatCompletionOptionContainsMessages() private ChatCompletionsClient CreateChatCompletionClient() { var apiKey = Environment.GetEnvironmentVariable("GH_API_KEY") ?? throw new Exception("Please set GH_API_KEY environment variable."); - var endpoint = "https://models.inference.ai.azure.com"; + var endpoint = "https://models.github.ai/inference"; return new ChatCompletionsClient(new Uri(endpoint), new Azure.AzureKeyCredential(apiKey)); } diff --git a/dotnet/test/AutoGen.OpenAI.Tests/ApprovalTests/OpenAIMessageTests.BasicMessageTest.approved.txt b/dotnet/test/AutoGen.OpenAI.Tests/ApprovalTests/OpenAIMessageTests.BasicMessageTest.approved.txt index 5b113c3f65ab..78c374f63d8e 100644 --- a/dotnet/test/AutoGen.OpenAI.Tests/ApprovalTests/OpenAIMessageTests.BasicMessageTest.approved.txt +++ b/dotnet/test/AutoGen.OpenAI.Tests/ApprovalTests/OpenAIMessageTests.BasicMessageTest.approved.txt @@ -12,6 +12,12 @@ "ImageUri": null, "ImageBytes": null, "ImageBytesMediaType": null, + "InputAudioBytes": null, + "InputAudioFormat": null, + "FileId": null, + "FileBytes": null, + "FileBytesMediaType": null, + "Filename": null, "ImageDetailLevel": null, "Refusal": null } @@ -31,6 +37,12 @@ "ImageUri": null, "ImageBytes": null, "ImageBytesMediaType": null, + "InputAudioBytes": null, + "InputAudioFormat": null, + "FileId": null, + "FileBytes": null, + "FileBytesMediaType": null, + "Filename": null, "ImageDetailLevel": null, "Refusal": null } @@ -57,6 +69,12 @@ "ImageUri": null, "ImageBytes": null, "ImageBytesMediaType": null, + "InputAudioBytes": null, + "InputAudioFormat": null, + "FileId": null, + "FileBytes": null, + "FileBytesMediaType": null, + "Filename": null, "ImageDetailLevel": null, "Refusal": null } @@ -78,6 +96,12 @@ "ImageUri": "https://example.com/image.png", "ImageBytes": null, "ImageBytesMediaType": null, + "InputAudioBytes": null, + "InputAudioFormat": null, + "FileId": null, + "FileBytes": null, + "FileBytesMediaType": null, + "Filename": null, "ImageDetailLevel": null, "Refusal": null } @@ -104,6 +128,12 @@ "ImageUri": null, "ImageBytes": null, "ImageBytesMediaType": null, + "InputAudioBytes": null, + "InputAudioFormat": null, + "FileId": null, + "FileBytes": null, + "FileBytesMediaType": null, + "Filename": null, "ImageDetailLevel": null, "Refusal": null }, @@ -113,6 +143,12 @@ "ImageUri": "https://example.com/image.png", "ImageBytes": null, "ImageBytesMediaType": null, + "InputAudioBytes": null, + "InputAudioFormat": null, + "FileId": null, + "FileBytes": null, + "FileBytesMediaType": null, + "Filename": null, "ImageDetailLevel": null, "Refusal": null } diff --git a/dotnet/test/AutoGen.OpenAI.Tests/ApprovalTests/OpenAIMessageTests.BasicMessageTest.received.txt b/dotnet/test/AutoGen.OpenAI.Tests/ApprovalTests/OpenAIMessageTests.BasicMessageTest.received.txt new file mode 100644 index 000000000000..78c374f63d8e --- /dev/null +++ b/dotnet/test/AutoGen.OpenAI.Tests/ApprovalTests/OpenAIMessageTests.BasicMessageTest.received.txt @@ -0,0 +1,260 @@ +[ + { + "OriginalMessage": "TextMessage(system, You are a helpful AI assistant, )", + "ConvertedMessages": [ + { + "Name": null, + "Role": "system", + "Content": [ + { + "Kind": 0, + "Text": "You are a helpful AI assistant", + "ImageUri": null, + "ImageBytes": null, + "ImageBytesMediaType": null, + "InputAudioBytes": null, + "InputAudioFormat": null, + "FileId": null, + "FileBytes": null, + "FileBytesMediaType": null, + "Filename": null, + "ImageDetailLevel": null, + "Refusal": null + } + ] + } + ] + }, + { + "OriginalMessage": "TextMessage(user, Hello, user)", + "ConvertedMessages": [ + { + "Role": "user", + "Content": [ + { + "Kind": 0, + "Text": "Hello", + "ImageUri": null, + "ImageBytes": null, + "ImageBytesMediaType": null, + "InputAudioBytes": null, + "InputAudioFormat": null, + "FileId": null, + "FileBytes": null, + "FileBytesMediaType": null, + "Filename": null, + "ImageDetailLevel": null, + "Refusal": null + } + ], + "Name": "user", + "MultiModaItem": [ + { + "Type": "Text", + "Text": "Hello" + } + ] + } + ] + }, + { + "OriginalMessage": "TextMessage(assistant, How can I help you?, assistant)", + "ConvertedMessages": [ + { + "Role": "assistant", + "Content": [ + { + "Kind": 0, + "Text": "How can I help you?", + "ImageUri": null, + "ImageBytes": null, + "ImageBytesMediaType": null, + "InputAudioBytes": null, + "InputAudioFormat": null, + "FileId": null, + "FileBytes": null, + "FileBytesMediaType": null, + "Filename": null, + "ImageDetailLevel": null, + "Refusal": null + } + ], + "Name": "assistant", + "TooCall": [] + } + ] + }, + { + "OriginalMessage": "ImageMessage(user, https://example.com/image.png, user)", + "ConvertedMessages": [ + { + "Role": "user", + "Content": [ + { + "Kind": 2, + "Text": null, + "ImageUri": "https://example.com/image.png", + "ImageBytes": null, + "ImageBytesMediaType": null, + "InputAudioBytes": null, + "InputAudioFormat": null, + "FileId": null, + "FileBytes": null, + "FileBytesMediaType": null, + "Filename": null, + "ImageDetailLevel": null, + "Refusal": null + } + ], + "Name": "user", + "MultiModaItem": [ + { + "Type": "Image", + "ImageUrl": "https://example.com/image.png" + } + ] + } + ] + }, + { + "OriginalMessage": "MultiModalMessage(assistant, user)\n\tTextMessage(user, Hello, user)\n\tImageMessage(user, https://example.com/image.png, user)", + "ConvertedMessages": [ + { + "Role": "user", + "Content": [ + { + "Kind": 0, + "Text": "Hello", + "ImageUri": null, + "ImageBytes": null, + "ImageBytesMediaType": null, + "InputAudioBytes": null, + "InputAudioFormat": null, + "FileId": null, + "FileBytes": null, + "FileBytesMediaType": null, + "Filename": null, + "ImageDetailLevel": null, + "Refusal": null + }, + { + "Kind": 2, + "Text": null, + "ImageUri": "https://example.com/image.png", + "ImageBytes": null, + "ImageBytesMediaType": null, + "InputAudioBytes": null, + "InputAudioFormat": null, + "FileId": null, + "FileBytes": null, + "FileBytesMediaType": null, + "Filename": null, + "ImageDetailLevel": null, + "Refusal": null + } + ], + "Name": "user", + "MultiModaItem": [ + { + "Type": "Text", + "Text": "Hello" + }, + { + "Type": "Image", + "ImageUrl": "https://example.com/image.png" + } + ] + } + ] + }, + { + "OriginalMessage": "ToolCallMessage(assistant)\n\tToolCall(test, test, )", + "ConvertedMessages": [ + { + "Role": "assistant", + "Content": [], + "Name": null, + "TooCall": [ + { + "Type": "Function", + "Name": "test", + "Arguments": "dGVzdA==", + "Id": "test" + } + ] + } + ] + }, + { + "OriginalMessage": "ToolCallResultMessage(user)\n\tToolCall(test, test, result)", + "ConvertedMessages": [ + { + "Role": "tool", + "Content": "result", + "ToolCallId": "test" + } + ] + }, + { + "OriginalMessage": "ToolCallResultMessage(user)\n\tToolCall(result, test, test)\n\tToolCall(result, test, test)", + "ConvertedMessages": [ + { + "Role": "tool", + "Content": "test", + "ToolCallId": "result_0" + }, + { + "Role": "tool", + "Content": "test", + "ToolCallId": "result_1" + } + ] + }, + { + "OriginalMessage": "ToolCallMessage(assistant)\n\tToolCall(test, test, )\n\tToolCall(test, test, )", + "ConvertedMessages": [ + { + "Role": "assistant", + "Content": [], + "Name": null, + "TooCall": [ + { + "Type": "Function", + "Name": "test", + "Arguments": "dGVzdA==", + "Id": "test_0" + }, + { + "Type": "Function", + "Name": "test", + "Arguments": "dGVzdA==", + "Id": "test_1" + } + ] + } + ] + }, + { + "OriginalMessage": "AggregateMessage(assistant)\n\tToolCallMessage(assistant)\n\tToolCall(test, test, )\n\tToolCallResultMessage(assistant)\n\tToolCall(test, test, result)", + "ConvertedMessages": [ + { + "Role": "assistant", + "Content": [], + "Name": null, + "TooCall": [ + { + "Type": "Function", + "Name": "test", + "Arguments": "dGVzdA==", + "Id": "test" + } + ] + }, + { + "Role": "tool", + "Content": "result", + "ToolCallId": "test" + } + ] + } +] \ No newline at end of file diff --git a/dotnet/test/AutoGen.Tests/Function/FunctionTests.cs b/dotnet/test/AutoGen.Tests/Function/FunctionTests.cs index adb97e1c5327..b8024e0a360e 100644 --- a/dotnet/test/AutoGen.Tests/Function/FunctionTests.cs +++ b/dotnet/test/AutoGen.Tests/Function/FunctionTests.cs @@ -60,7 +60,7 @@ public async Task CreateGetWeatherFunctionFromAIFunctionFactoryAsync() GetWeatherAsyncStatic, ]; - var functionContracts = availableDelegates.Select(function => (FunctionContract)AIFunctionFactory.Create(function).Metadata).ToList(); + var functionContracts = availableDelegates.Select(function => (FunctionContract)AIFunctionFactory.Create(function)).ToList(); // Verify the function contracts functionContracts.Should().HaveCount(4); diff --git a/dotnet/test/AutoGen.Tests/Orchestrator/RolePlayOrchestratorTests.cs b/dotnet/test/AutoGen.Tests/Orchestrator/RolePlayOrchestratorTests.cs index 256d704cc4c4..06b2b6d5f2bc 100644 --- a/dotnet/test/AutoGen.Tests/Orchestrator/RolePlayOrchestratorTests.cs +++ b/dotnet/test/AutoGen.Tests/Orchestrator/RolePlayOrchestratorTests.cs @@ -18,7 +18,6 @@ using AutoGen.OpenAI; using AutoGen.OpenAI.Extension; using Azure.AI.Inference; -using Azure.AI.OpenAI; using FluentAssertions; using Moq; using OpenAI; @@ -217,21 +216,6 @@ public async Task ItUseCandidatesFromWorflowAsync() speaker.Should().Be(bob); } - [ApiKeyFact("AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_DEPLOY_NAME")] - public async Task GPT_3_5_CoderReviewerRunnerTestAsync() - { - var endpoint = Environment.GetEnvironmentVariable("AZURE_OPENAI_ENDPOINT") ?? throw new Exception("Please set AZURE_OPENAI_ENDPOINT environment variable."); - var key = Environment.GetEnvironmentVariable("AZURE_OPENAI_API_KEY") ?? throw new Exception("Please set AZURE_OPENAI_API_KEY environment variable."); - var deployName = Environment.GetEnvironmentVariable("AZURE_OPENAI_DEPLOY_NAME") ?? throw new Exception("Please set AZURE_OPENAI_DEPLOY_NAME environment variable."); - var openaiClient = new AzureOpenAIClient(new Uri(endpoint), new System.ClientModel.ApiKeyCredential(key)); - var openAIChatAgent = new OpenAIChatAgent( - chatClient: openaiClient.GetChatClient(deployName), - name: "assistant") - .RegisterMessageConnector(); - - await CoderReviewerRunnerTestAsync(openAIChatAgent); - } - [ApiKeyFact("OPENAI_API_KEY")] public async Task GPT_4o_CoderReviewerRunnerTestAsync() { @@ -308,7 +292,7 @@ public async Task Mistra_7b_CoderReviewerRunnerTestAsync() public async Task LLaMA_3_1_CoderReviewerRunnerTestAsync() { var apiKey = Environment.GetEnvironmentVariable("GH_API_KEY") ?? throw new InvalidOperationException("GH_API_KEY is not set."); - var endPoint = "https://models.inference.ai.azure.com"; + var endPoint = "https://models.github.ai/inference"; var chatCompletionClient = new ChatCompletionsClient(new Uri(endPoint), new Azure.AzureKeyCredential(apiKey)); var agent = new ChatCompletionsClientAgent( diff --git a/dotnet/test/AutoGen.Tests/SingleAgentTest.cs b/dotnet/test/AutoGen.Tests/SingleAgentTest.cs index 8611714c351d..2f014236c403 100644 --- a/dotnet/test/AutoGen.Tests/SingleAgentTest.cs +++ b/dotnet/test/AutoGen.Tests/SingleAgentTest.cs @@ -20,11 +20,11 @@ public SingleAgentTest(ITestOutputHelper output) _output = output; } - private ILLMConfig CreateAzureOpenAIGPT35TurboConfig() + private ILLMConfig CreateAzureOpenAIGPT4oMiniConfig() { var key = Environment.GetEnvironmentVariable("AZURE_OPENAI_API_KEY") ?? throw new ArgumentException("AZURE_OPENAI_API_KEY is not set"); var endpoint = Environment.GetEnvironmentVariable("AZURE_OPENAI_ENDPOINT") ?? throw new ArgumentException("AZURE_OPENAI_ENDPOINT is not set"); - var deployName = Environment.GetEnvironmentVariable("AZURE_OPENAI_DEPLOY_NAME") ?? throw new ArgumentException("AZURE_OPENAI_DEPLOY_NAME is not set"); + var deployName = "gpt-4o-mini"; return new AzureOpenAIConfig(endpoint, deployName, key); } @@ -37,7 +37,7 @@ private ILLMConfig CreateOpenAIGPT4VisionConfig() [ApiKeyFact("AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_DEPLOY_NAME")] public async Task AssistantAgentFunctionCallTestAsync() { - var config = this.CreateAzureOpenAIGPT35TurboConfig(); + var config = this.CreateAzureOpenAIGPT4oMiniConfig(); var llmConfig = new ConversableAgentConfig { @@ -77,7 +77,7 @@ public async Task AssistantAgentDefaultReplyTestAsync() [ApiKeyFact("AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_DEPLOY_NAME")] public async Task AssistantAgentFunctionCallSelfExecutionTestAsync() { - var config = this.CreateAzureOpenAIGPT35TurboConfig(); + var config = this.CreateAzureOpenAIGPT4oMiniConfig(); var llmConfig = new ConversableAgentConfig { FunctionContracts = new[] diff --git a/dotnet/test/AutoGen.Tests/TwoAgentTest.cs b/dotnet/test/AutoGen.Tests/TwoAgentTest.cs index 8ba4a3fedbd8..ed24dabef3bf 100644 --- a/dotnet/test/AutoGen.Tests/TwoAgentTest.cs +++ b/dotnet/test/AutoGen.Tests/TwoAgentTest.cs @@ -32,7 +32,7 @@ public async Task TwoAgentWeatherChatTestAsync() { var key = Environment.GetEnvironmentVariable("AZURE_OPENAI_API_KEY") ?? throw new ArgumentException("AZURE_OPENAI_API_KEY is not set"); var endpoint = Environment.GetEnvironmentVariable("AZURE_OPENAI_ENDPOINT") ?? throw new ArgumentException("AZURE_OPENAI_ENDPOINT is not set"); - var deploymentName = Environment.GetEnvironmentVariable("AZURE_OPENAI_DEPLOY_NAME") ?? throw new ArgumentException("AZURE_OPENAI_DEPLOY_NAME is not set"); + var deploymentName = "gpt-4o-mini"; var config = new AzureOpenAIConfig(endpoint, deploymentName, key); var assistant = new AssistantAgent( diff --git a/python/.gitignore b/python/.gitignore index d7d27c394883..e5b2be591395 100644 --- a/python/.gitignore +++ b/python/.gitignore @@ -174,3 +174,6 @@ tmp_code_*.py # .NET Development settings appsettings.Development.json + +# Documentation reference files +docs/src/reference \ No newline at end of file diff --git a/python/README.md b/python/README.md index f163de315375..df1f24288093 100644 --- a/python/README.md +++ b/python/README.md @@ -1,15 +1,21 @@ -# AutoGen Python packages +# AutoGen Python Development Guide -[![0.4 Docs](https://img.shields.io/badge/Docs-0.4-blue)](https://microsoft.github.io/autogen/dev/) +[![Docs (dev)](https://img.shields.io/badge/Docs-dev-blue)](https://microsoft.github.io/autogen/dev/) +[![Docs (latest release)](https://img.shields.io/badge/Docs-latest%20release-blue)](https://microsoft.github.io/autogen/dev/) [![PyPi autogen-core](https://img.shields.io/badge/PyPi-autogen--core-blue?logo=pypi)](https://pypi.org/project/autogen-core/) [![PyPi autogen-agentchat](https://img.shields.io/badge/PyPi-autogen--agentchat-blue?logo=pypi)](https://pypi.org/project/autogen-agentchat/) [![PyPi autogen-ext](https://img.shields.io/badge/PyPi-autogen--ext-blue?logo=pypi)](https://pypi.org/project/autogen-ext/) -This directory works as a single `uv` workspace containing all project packages. See [`packages`](./packages/) to discover all project packages. +This directory works as a single `uv` workspace containing all project packages, including: + +- `packages/autogen-core`: interface definitions and reference implementations of agent runtime, model, tool, workbench, memory, tracing. +- `packages/autogen-agentchat`: single and multi-agent workflows built on top of `autogen-core`. +- `packages/autogen-ext`: implementations for ecosystem integrations. For example, `autogen-ext[openai]` provides the OpenAI model client. +- `packages/autogen-studio`: a web-based IDE for building and running AutoGen agents. ## Migrating from 0.2.x? Please refer to the [migration guide](./migration_guide.md) for how to migrate your code from 0.2.x to 0.4.x. -## Development +## Quick Start **TL;DR**, run all checks with: @@ -19,20 +25,19 @@ source .venv/bin/activate poe check ``` -### Setup +## Setup `uv` is a package manager that assists in creating the necessary environment and installing packages to run AutoGen. - [Install `uv`](https://docs.astral.sh/uv/getting-started/installation/). -**Note:** To prevent incompatibilities between versions the same UV version as is running in CI should be used. Check the version in CI by looking the `setup-uv` action, [here](https://github.com/microsoft/autogen/blob/main/.github/workflows/checks.yml#L40) for example. +To upgrade `uv` to the latest version, run: -For example, to change your version to `0.5.18`, run: ```sh -uv self update 0.5.18 +uv self update ``` -### Virtual Environment +## Virtual Environment During development, you may need to test changes made to any of the packages.\ To do so, create a virtual environment where the AutoGen packages are installed based on the current state of the directory.\ @@ -46,7 +51,7 @@ source .venv/bin/activate - `uv sync --all-extras` will create a `.venv` directory at the current level and install packages from the current directory along with any other dependencies. The `all-extras` flag adds optional dependencies. - `source .venv/bin/activate` activates the virtual environment. -### Common Tasks +## Common Tasks To create a pull request (PR), ensure the following checks are met. You can run each check individually: @@ -55,16 +60,19 @@ To create a pull request (PR), ensure the following checks are met. You can run - Test: `poe test` - Mypy: `poe mypy` - Pyright: `poe pyright` -- Build docs: `poe --directory ./packages/autogen-core/ docs-build` -- Auto rebuild+serve docs: `poe --directory ./packages/autogen-core/ docs-serve` +- Build docs: `poe docs-build` +- Check docs: `poe docs-check` +- Clean docs: `poe docs-clean` +- Check code blocks in API references: `poe docs-check-examples` +- Auto rebuild+serve docs: `poe docs-serve` - Check samples in `python/samples`: `poe samples-code-check` -Alternatively, you can run all the checks with: + Alternatively, you can run all the checks with: - `poe check` > [!NOTE] > These need to be run in the virtual environment. -### Syncing Dependencies +## Syncing Dependencies When you pull new changes, you may need to update the dependencies. To do so, first make sure you are in the virtual environment, and then in the `python` directory, run: @@ -75,7 +83,135 @@ uv sync --all-extras This will update the dependencies in the virtual environment. -### Creating a New Package +## Building Documentation + +The documentation source directory is located at `docs/src/`. + +To build the documentation, run this from the root of the Python directory: + +```sh +poe docs-build +``` + +To serve the documentation locally, run: + +```sh +poe docs-serve +``` + +When you make changes to the doc strings or add new modules, you may need to +refresh the API references in the documentation by first cleaning the docs and +then building them again: + +```sh +poe docs-clean # This will remove the build directory and the reference directory +poe docs-build # This will rebuild the documentation from scratch +``` + +## Writing Documentation + +When you add a new public class or function, you should always add a docstring +to it. The docstring should follow the +[Google style](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) layout +and the Sphinx RST format for Python docstrings. + +The docstring for a public class or function should include: + +- A short description of the class or function at the beginning immediately after the `"""`. +- A longer description if necessary, explaining the purpose and usage. +- A list of arguments with their types and descriptions, using the `Args` section. + Each argument should be listed with its name, type, and a brief description. +- A description of the return value and its type, using the `Returns` section. + If the function does not return anything, you can omit this section. +- A list of exceptions that the function may raise, with descriptions, + using the `Raises` section. This is optional but recommended if the function can raise exceptions that users should be aware of. +- Examples of how to use the class or function, using the `Examples` section, + and formatted using `.. code-block:: python` directive. Optionally, also include the output of the example using + `.. code-block:: text` directive. + +Here is an example of a docstring for `McpWorkbench` class: + +```python +class McpWorkbench(Workbench, Component[McpWorkbenchConfig]): + """A workbench that wraps an MCP server and provides an interface + to list and call tools provided by the server. + + This workbench should be used as a context manager to ensure proper + initialization and cleanup of the underlying MCP session. + + Args: + server_params (McpServerParams): The parameters to connect to the MCP server. + This can be either a :class:`StdioServerParams` or :class:`SseServerParams`. + tool_overrides (Optional[Dict[str, ToolOverride]]): Optional mapping of original tool + names to override configurations for name and/or description. This allows + customizing how server tools appear to consumers while maintaining the underlying + tool functionality. + + Raises: + ValueError: If there are conflicts in tool override names. + + Examples: + + Here is a simple example of how to use the workbench with a `mcp-server-fetch` server: + + .. code-block:: python + + import asyncio + + from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams + + + async def main() -> None: + params = StdioServerParams( + command="uvx", + args=["mcp-server-fetch"], + read_timeout_seconds=60, + ) + + # You can also use `start()` and `stop()` to manage the session. + async with McpWorkbench(server_params=params) as workbench: + tools = await workbench.list_tools() + print(tools) + result = await workbench.call_tool(tools[0]["name"], {"url": "https://github.com/"}) + print(result) + + + asyncio.run(main()) +``` + +The code blocks with `.. code-block:: python` is checked by the `docs-check-examples` task using Pyright, +so make sure the code is valid. Running the code as a script and checking it using `pyright` +is a good way to ensure the code examples are correct. + +When you reference a class, method, or function in the docstring, you should always +use the `:class:`, `:meth:`, or `:func:` directive to create a link to the class or function. +Always use the fully qualified name of the class or function, including the package name, but +prefix it with a `~` for shorter rendering in the documentation. +For example, if you are referencing the `AssistantAgent` class in the `autogen-agentchat` package, +you should write it as `:class:~autogen_agentchat.AssistantAgent`. + +For a public data class, including those that are Pydantic models, you should also include docstrings +for each field in the class. + +## Writing Tests + +When you add a new public class or function, you should also always add tests for it. +We track test coverage and aim for not reducing the coverage percentage with new changes. + +We use `pytest` for testing, and you should always use fixtures to set up the test dependencies. + +Use mock objects to simulate dependencies and avoid making real API calls or database queries in tests. +See existing tests for examples of how to use fixtures and mocks. + +For model clients, use `autogen_ext.models.replay.ReplayChatCompletionClient` as a +drop-in replacement for the model client to simulate responses without making real API calls. + +When certain tests requires interaction with actual model APIs or other external services, +you should configure the tests to be skipped if the required services are not available. +For example, if you are testing a model client that requires an OpenAI API key, +you can use the `pytest.mark.skipif` decorator to skip the test if the environment variable for the API key is not set. + +## Creating a New Package To create a new package, similar to `autogen-core` or `autogen-chat`, use the following: diff --git a/python/docs/README.md b/python/docs/README.md new file mode 100644 index 000000000000..7b813527a789 --- /dev/null +++ b/python/docs/README.md @@ -0,0 +1,29 @@ +## Building the AutoGen Documentation + +AutoGen documentation is based on the sphinx documentation system and uses the myst-parser to render markdown files. It uses the [pydata-sphinx-theme](https://pydata-sphinx-theme.readthedocs.io/en/latest/) to style the documentation. + +### Prerequisites + +Ensure you have all of the dev dependencies for the `autogen-core` package installed. You can install them by running the following command from the root of the python repository: + +```bash +uv sync +source .venv/bin/activate +``` + +## Building Docs + +To build the documentation, run the following command from the root of the python directory: + +```bash +poe docs-build +``` + +To serve the documentation locally, run the following command from the root of the python directory: + +```bash +poe docs-serve +``` + +[!NOTE] +Sphinx will only rebuild files that have changed since the last build. If you want to force a full rebuild, you can delete the `./docs/build` directory before running the `docs-build` command. diff --git a/python/packages/autogen-core/docs/drawio/agent-lifecycle.drawio b/python/docs/drawio/agent-lifecycle.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/agent-lifecycle.drawio rename to python/docs/drawio/agent-lifecycle.drawio diff --git a/python/packages/autogen-core/docs/drawio/agentchat-team.drawio b/python/docs/drawio/agentchat-team.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/agentchat-team.drawio rename to python/docs/drawio/agentchat-team.drawio diff --git a/python/packages/autogen-core/docs/drawio/application-stack.drawio b/python/docs/drawio/application-stack.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/application-stack.drawio rename to python/docs/drawio/application-stack.drawio diff --git a/python/packages/autogen-core/docs/drawio/architecture-distributed.drawio b/python/docs/drawio/architecture-distributed.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/architecture-distributed.drawio rename to python/docs/drawio/architecture-distributed.drawio diff --git a/python/packages/autogen-core/docs/drawio/architecture-standalone.drawio b/python/docs/drawio/architecture-standalone.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/architecture-standalone.drawio rename to python/docs/drawio/architecture-standalone.drawio diff --git a/python/packages/autogen-core/docs/drawio/assistant-agent.drawio b/python/docs/drawio/assistant-agent.drawio similarity index 64% rename from python/packages/autogen-core/docs/drawio/assistant-agent.drawio rename to python/docs/drawio/assistant-agent.drawio index 709bb3727b20..377bfbdf62d9 100644 --- a/python/packages/autogen-core/docs/drawio/assistant-agent.drawio +++ b/python/docs/drawio/assistant-agent.drawio @@ -1,202 +1,222 @@ - + - + - - + + - - + + - + - + - + - + - + - + - + - + - + - - + + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - - + + - + - + - - + + - - + + - - + + - - + + - - + + - - + + - - + + - + - + - + - + - + - + - + - + - + - + - - + + - + - - + + - + - + - + - + + + + - + - + - + + + + + + + + + + + + + + + + + + diff --git a/python/packages/autogen-core/docs/drawio/code-gen-example.drawio b/python/docs/drawio/code-gen-example.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/code-gen-example.drawio rename to python/docs/drawio/code-gen-example.drawio diff --git a/python/packages/autogen-core/docs/drawio/coder-reviewer.drawio b/python/docs/drawio/coder-reviewer.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/coder-reviewer.drawio rename to python/docs/drawio/coder-reviewer.drawio diff --git a/python/packages/autogen-core/docs/drawio/groupchat.drawio b/python/docs/drawio/groupchat.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/groupchat.drawio rename to python/docs/drawio/groupchat.drawio diff --git a/python/packages/autogen-core/docs/drawio/handoffs.drawio b/python/docs/drawio/handoffs.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/handoffs.drawio rename to python/docs/drawio/handoffs.drawio diff --git a/python/packages/autogen-core/docs/drawio/human-in-the-loop-termination.drawio b/python/docs/drawio/human-in-the-loop-termination.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/human-in-the-loop-termination.drawio rename to python/docs/drawio/human-in-the-loop-termination.drawio diff --git a/python/packages/autogen-core/docs/drawio/human-in-the-loop-user-proxy.drawio b/python/docs/drawio/human-in-the-loop-user-proxy.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/human-in-the-loop-user-proxy.drawio rename to python/docs/drawio/human-in-the-loop-user-proxy.drawio diff --git a/python/packages/autogen-core/docs/drawio/selector-group-chat.drawio b/python/docs/drawio/selector-group-chat.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/selector-group-chat.drawio rename to python/docs/drawio/selector-group-chat.drawio diff --git a/python/packages/autogen-core/docs/drawio/sequential-workflow.drawio b/python/docs/drawio/sequential-workflow.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/sequential-workflow.drawio rename to python/docs/drawio/sequential-workflow.drawio diff --git a/python/packages/autogen-core/docs/drawio/subscription.drawio b/python/docs/drawio/subscription.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/subscription.drawio rename to python/docs/drawio/subscription.drawio diff --git a/python/packages/autogen-core/docs/drawio/swarm_customer_support.drawio b/python/docs/drawio/swarm_customer_support.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/swarm_customer_support.drawio rename to python/docs/drawio/swarm_customer_support.drawio diff --git a/python/packages/autogen-core/docs/drawio/swarm_stock_research.drawio b/python/docs/drawio/swarm_stock_research.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/swarm_stock_research.drawio rename to python/docs/drawio/swarm_stock_research.drawio diff --git a/python/packages/autogen-core/docs/drawio/type-subscription.drawio b/python/docs/drawio/type-subscription.drawio similarity index 100% rename from python/packages/autogen-core/docs/drawio/type-subscription.drawio rename to python/docs/drawio/type-subscription.drawio diff --git a/python/packages/autogen-core/docs/redirects/redirect_template.html b/python/docs/redirects/redirect_template.html similarity index 100% rename from python/packages/autogen-core/docs/redirects/redirect_template.html rename to python/docs/redirects/redirect_template.html diff --git a/python/packages/autogen-core/docs/redirects/redirect_urls.txt b/python/docs/redirects/redirect_urls.txt similarity index 100% rename from python/packages/autogen-core/docs/redirects/redirect_urls.txt rename to python/docs/redirects/redirect_urls.txt diff --git a/python/packages/autogen-core/docs/redirects/redirects.py b/python/docs/redirects/redirects.py similarity index 100% rename from python/packages/autogen-core/docs/redirects/redirects.py rename to python/docs/redirects/redirects.py diff --git a/python/packages/autogen-core/docs/src/_apidoc_templates/module.rst.jinja b/python/docs/src/_apidoc_templates/module.rst.jinja similarity index 100% rename from python/packages/autogen-core/docs/src/_apidoc_templates/module.rst.jinja rename to python/docs/src/_apidoc_templates/module.rst.jinja diff --git a/python/packages/autogen-core/docs/src/_apidoc_templates/package.rst.jinja b/python/docs/src/_apidoc_templates/package.rst.jinja similarity index 100% rename from python/packages/autogen-core/docs/src/_apidoc_templates/package.rst.jinja rename to python/docs/src/_apidoc_templates/package.rst.jinja diff --git a/python/packages/autogen-core/docs/src/_extension/code_lint.py b/python/docs/src/_extension/code_lint.py similarity index 100% rename from python/packages/autogen-core/docs/src/_extension/code_lint.py rename to python/docs/src/_extension/code_lint.py diff --git a/python/packages/autogen-core/docs/src/_extension/gallery_directive.py b/python/docs/src/_extension/gallery_directive.py similarity index 100% rename from python/packages/autogen-core/docs/src/_extension/gallery_directive.py rename to python/docs/src/_extension/gallery_directive.py diff --git a/python/packages/autogen-core/docs/src/_static/banner-override.js b/python/docs/src/_static/banner-override.js similarity index 100% rename from python/packages/autogen-core/docs/src/_static/banner-override.js rename to python/docs/src/_static/banner-override.js diff --git a/python/packages/autogen-core/docs/src/_static/custom-icon.js b/python/docs/src/_static/custom-icon.js similarity index 100% rename from python/packages/autogen-core/docs/src/_static/custom-icon.js rename to python/docs/src/_static/custom-icon.js diff --git a/python/packages/autogen-core/docs/src/_static/custom.css b/python/docs/src/_static/custom.css similarity index 91% rename from python/packages/autogen-core/docs/src/_static/custom.css rename to python/docs/src/_static/custom.css index ce750411ba63..ae67f8c882e6 100644 --- a/python/packages/autogen-core/docs/src/_static/custom.css +++ b/python/docs/src/_static/custom.css @@ -32,6 +32,17 @@ html[data-theme="dark"] { color: white; text-shadow: 0.5px 0 0 currentColor; } + +/* Adding header icon hover and focus effects */ +.bd-header a:focus-visible { + color: var(--pst-color-secondary) !important; + text-decoration: underline !important; + text-shadow: 0.5px 0 0 currentColor; + transform: scale(1.05); + transition: all 0.2s ease-in-out; + outline: none; +} + nav.bd-links .current>a { box-shadow: inset 1px 0 0 var(--pst-color-primary); } diff --git a/python/packages/autogen-core/docs/src/_static/custom.js b/python/docs/src/_static/custom.js similarity index 100% rename from python/packages/autogen-core/docs/src/_static/custom.js rename to python/docs/src/_static/custom.js diff --git a/python/packages/autogen-core/docs/src/_static/images/logo/favicon-16x16.png b/python/docs/src/_static/images/logo/favicon-16x16.png similarity index 100% rename from python/packages/autogen-core/docs/src/_static/images/logo/favicon-16x16.png rename to python/docs/src/_static/images/logo/favicon-16x16.png diff --git a/python/packages/autogen-core/docs/src/_static/images/logo/favicon-32x32.png b/python/docs/src/_static/images/logo/favicon-32x32.png similarity index 100% rename from python/packages/autogen-core/docs/src/_static/images/logo/favicon-32x32.png rename to python/docs/src/_static/images/logo/favicon-32x32.png diff --git a/python/packages/autogen-core/docs/src/_static/images/logo/favicon-512x512.png b/python/docs/src/_static/images/logo/favicon-512x512.png similarity index 100% rename from python/packages/autogen-core/docs/src/_static/images/logo/favicon-512x512.png rename to python/docs/src/_static/images/logo/favicon-512x512.png diff --git a/python/packages/autogen-core/docs/src/_static/images/logo/favicon.ico b/python/docs/src/_static/images/logo/favicon.ico similarity index 100% rename from python/packages/autogen-core/docs/src/_static/images/logo/favicon.ico rename to python/docs/src/_static/images/logo/favicon.ico diff --git a/python/packages/autogen-core/docs/src/_static/images/logo/logo.svg b/python/docs/src/_static/images/logo/logo.svg similarity index 100% rename from python/packages/autogen-core/docs/src/_static/images/logo/logo.svg rename to python/docs/src/_static/images/logo/logo.svg diff --git a/python/packages/autogen-core/docs/src/_static/switcher.json b/python/docs/src/_static/switcher.json similarity index 100% rename from python/packages/autogen-core/docs/src/_static/switcher.json rename to python/docs/src/_static/switcher.json diff --git a/python/packages/autogen-core/docs/src/_templates/edit-this-page.html b/python/docs/src/_templates/edit-this-page.html similarity index 100% rename from python/packages/autogen-core/docs/src/_templates/edit-this-page.html rename to python/docs/src/_templates/edit-this-page.html diff --git a/python/packages/autogen-core/docs/src/_templates/footer-middle-links.html b/python/docs/src/_templates/footer-middle-links.html similarity index 100% rename from python/packages/autogen-core/docs/src/_templates/footer-middle-links.html rename to python/docs/src/_templates/footer-middle-links.html diff --git a/python/packages/autogen-core/docs/src/_templates/sidebar-nav-bs-agentchat.html b/python/docs/src/_templates/sidebar-nav-bs-agentchat.html similarity index 100% rename from python/packages/autogen-core/docs/src/_templates/sidebar-nav-bs-agentchat.html rename to python/docs/src/_templates/sidebar-nav-bs-agentchat.html diff --git a/python/packages/autogen-core/docs/src/_templates/sidebar-nav-bs-core.html b/python/docs/src/_templates/sidebar-nav-bs-core.html similarity index 100% rename from python/packages/autogen-core/docs/src/_templates/sidebar-nav-bs-core.html rename to python/docs/src/_templates/sidebar-nav-bs-core.html diff --git a/python/packages/autogen-core/docs/src/_templates/sidebar-nav-bs-extensions.html b/python/docs/src/_templates/sidebar-nav-bs-extensions.html similarity index 100% rename from python/packages/autogen-core/docs/src/_templates/sidebar-nav-bs-extensions.html rename to python/docs/src/_templates/sidebar-nav-bs-extensions.html diff --git a/python/packages/autogen-core/docs/src/_templates/sidebar-nav-bs-studio.html b/python/docs/src/_templates/sidebar-nav-bs-studio.html similarity index 100% rename from python/packages/autogen-core/docs/src/_templates/sidebar-nav-bs-studio.html rename to python/docs/src/_templates/sidebar-nav-bs-studio.html diff --git a/python/packages/autogen-core/docs/src/_templates/sidebar-nav-bs.html b/python/docs/src/_templates/sidebar-nav-bs.html similarity index 100% rename from python/packages/autogen-core/docs/src/_templates/sidebar-nav-bs.html rename to python/docs/src/_templates/sidebar-nav-bs.html diff --git a/python/packages/autogen-core/docs/src/_templates/theme-switcher.html b/python/docs/src/_templates/theme-switcher.html similarity index 100% rename from python/packages/autogen-core/docs/src/_templates/theme-switcher.html rename to python/docs/src/_templates/theme-switcher.html diff --git a/python/packages/autogen-core/docs/src/_templates/version-banner-override.html b/python/docs/src/_templates/version-banner-override.html similarity index 100% rename from python/packages/autogen-core/docs/src/_templates/version-banner-override.html rename to python/docs/src/_templates/version-banner-override.html diff --git a/python/packages/autogen-core/docs/src/conf.py b/python/docs/src/conf.py similarity index 82% rename from python/packages/autogen-core/docs/src/conf.py rename to python/docs/src/conf.py index ff5608a17d67..4c25c8e5fcd0 100644 --- a/python/packages/autogen-core/docs/src/conf.py +++ b/python/docs/src/conf.py @@ -8,6 +8,7 @@ from pathlib import Path import sys import os +import subprocess # -- Project information ----------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information @@ -149,7 +150,7 @@ "github_user": "microsoft", "github_repo": "autogen", "github_version": "main", - "doc_path": "python/packages/autogen-core/docs/src/", + "doc_path": "python/docs/src/", } autodoc_default_options = { @@ -182,6 +183,42 @@ } +def generate_api_reference() -> None: + """Generate API documentation before building.""" + reference_dir = Path(__file__).parent / "reference" + + # Only generate if reference directory doesn't exist + if reference_dir.exists(): + print("📁 Reference directory already exists, skipping API generation") + return + + script_path = Path(__file__).parent / "generate_api_reference.py" + if script_path.exists(): + print("🔄 Generating API documentation...") + try: + result = subprocess.run( + [sys.executable, str(script_path)], + cwd=script_path.parent, + capture_output=True, + text=True, + check=True + ) + print("✅ API documentation generated successfully") + # Print the output for visibility + if result.stdout: + for line in result.stdout.strip().split('\n'): + print(f" {line}") + except subprocess.CalledProcessError as e: + print(f"❌ Failed to generate API documentation: {e}") + if e.stdout: + print(f"stdout: {e.stdout}") + if e.stderr: + print(f"stderr: {e.stderr}") + # Don't fail the build, just warn + else: + print(f"⚠️ API documentation generator not found at {script_path}") + + def setup_to_main( app: Sphinx, pagename: str, templatename: str, context, doctree ) -> None: @@ -211,6 +248,9 @@ def setup(app: Sphinx) -> Dict[str, Any]: Returns: the 2 parallel parameters set to ``True``. """ + # Generate API documentation before building + app.connect("builder-inited", lambda app: generate_api_reference()) + app.connect("html-page-context", setup_to_main) # Adding here so it is inline and not in a separate file. diff --git a/python/docs/src/generate_api_reference.py b/python/docs/src/generate_api_reference.py new file mode 100644 index 000000000000..3b24957801f0 --- /dev/null +++ b/python/docs/src/generate_api_reference.py @@ -0,0 +1,305 @@ +#!/usr/bin/env python3 +""" +Script to automatically generate the API reference table of contents for AutoGen. + +This script scans all packages and their modules to generate the toctree entries +for the API documentation index.md file. +""" + +import os +from pathlib import Path +from typing import List, Dict, Set +import re + + +# Constants for package filtering and organization +DOCUMENTED_PACKAGES = ["autogen_core", "autogen_agentchat", "autogen_ext"] + +PACKAGE_SECTIONS = { + "autogen_agentchat": "AutoGen AgentChat", + "autogen_core": "AutoGen Core", + "autogen_ext": "AutoGen Extensions" +} + +# Exclusion patterns for submodules that are re-exported by parent modules +EXCLUSION_PATTERNS = [ + # task_centric_memory re-exports from memory_controller and utils + (r'^autogen_ext\.experimental\.task_centric_memory\.memory_controller$', + 'autogen_ext.experimental.task_centric_memory'), + # utils package re-exports from utils.apprentice and other utils submodules + (r'^autogen_ext\.experimental\.task_centric_memory\.utils\.apprentice$', + 'autogen_ext.experimental.task_centric_memory.utils'), + (r'^autogen_ext\.experimental\.task_centric_memory\.utils\.chat_completion_client_recorder$', + 'autogen_ext.experimental.task_centric_memory.utils'), + (r'^autogen_ext\.experimental\.task_centric_memory\.utils\.grader$', + 'autogen_ext.experimental.task_centric_memory.utils'), + (r'^autogen_ext\.experimental\.task_centric_memory\.utils\.page_logger$', + 'autogen_ext.experimental.task_centric_memory.utils'), + (r'^autogen_ext\.experimental\.task_centric_memory\.utils\.teachability$', + 'autogen_ext.experimental.task_centric_memory.utils'), +] + + +def is_private_module(module_parts: List[str]) -> bool: + """Check if any part of the module path indicates it's a private module.""" + return any(part.startswith('_') and part != '__init__' for part in module_parts) + + +def find_python_packages() -> List[Path]: + """Find documented Python packages in the workspace.""" + packages_dir = Path(__file__).parent.parent.parent / "packages" + python_packages = [] + + for package_dir in packages_dir.iterdir(): + if package_dir.is_dir(): + # Check if this package is in our documented packages list + package_name = package_dir.name.replace("-", "_") + if package_name in DOCUMENTED_PACKAGES: + src_dir = package_dir / "src" + if src_dir.exists(): + python_packages.append(src_dir) + + return python_packages + + +def get_module_hierarchy(package_root: Path) -> Dict[str, Set[str]]: + """Get the module hierarchy for a package, filtering only documented packages.""" + modules: Dict[str, Set[str]] = {} + + for root, dirs, files in os.walk(package_root): + # Skip __pycache__ and hidden directories + dirs[:] = [d for d in dirs if not d.startswith('__pycache__') and not d.startswith('.')] + + root_path = Path(root) + + # Process Python files (excluding private modules) + for file in files: + if file.endswith('.py') and file != '__init__.py' and not file.startswith('_'): + file_path = root_path / file + module_path = file_path.relative_to(package_root) + + # Convert file path to module name + module_parts = list(module_path.parts[:-1]) + [module_path.stem] + + if module_parts: + # Skip if any part of the module path is private + if is_private_module(module_parts): + continue + + module_name = '.'.join(module_parts) + package_name = module_parts[0] + + # Only include modules from documented packages + if package_name in DOCUMENTED_PACKAGES: + if package_name not in modules: + modules[package_name] = set() + + modules[package_name].add(module_name) + + # Also check for directories with __init__.py (packages, excluding private) + for dir_name in dirs: + if not dir_name.startswith('_'): # Skip private directories + dir_path = root_path / dir_name + if (dir_path / '__init__.py').exists(): + module_path = dir_path.relative_to(package_root) + module_parts = list(module_path.parts) + + if module_parts: + # Skip if any part of the module path is private + if is_private_module(module_parts): + continue + + module_name = '.'.join(module_parts) + package_name = module_parts[0] + + # Only include modules from documented packages + if package_name in DOCUMENTED_PACKAGES: + if package_name not in modules: + modules[package_name] = set() + + modules[package_name].add(module_name) + + return modules + + +def should_exclude_submodule(module_name: str, all_modules: Set[str]) -> bool: + """Check if a submodule should be excluded to avoid duplicate documentation.""" + for pattern, parent_module in EXCLUSION_PATTERNS: + if re.match(pattern, module_name) and parent_module in all_modules: + return True + + return False + + +def clean_rst_files(reference_dir: Path) -> None: + """Clean existing RST files to ensure fresh generation.""" + python_ref_dir = reference_dir / "python" + if python_ref_dir.exists(): + print("🧹 Cleaning existing .rst files...") + rst_files = list(python_ref_dir.glob("*.rst")) + for rst_file in rst_files: + rst_file.unlink() + print(f" Removed {len(rst_files)} existing .rst files") + + +def generate_rst_files(package_roots: List[Path], reference_dir: Path) -> Set[str]: + """Generate .rst files for all modules found in the packages.""" + python_ref_dir = reference_dir / "python" + python_ref_dir.mkdir(exist_ok=True, parents=True) + + # Clean existing RST files first + clean_rst_files(reference_dir) + + generated_files = set() + all_module_names = set() + + # First pass: collect all module names + for package_root in package_roots: + modules = get_module_hierarchy(package_root) + for package_name, module_set in modules.items(): + all_module_names.update(module_set) + + # Second pass: generate RST files, excluding problematic submodules + for package_root in package_roots: + modules = get_module_hierarchy(package_root) + + for package_name, module_set in modules.items(): + for module_name in module_set: + # Skip modules that would cause duplicate documentation + if should_exclude_submodule(module_name, all_module_names): + print(f" Skipping {module_name} (re-exported by parent)") + continue + + # Use the proper RST filename pattern (keep dots for submodules) + rst_filename = module_name + '.rst' + rst_path = python_ref_dir / rst_filename + + # Generate .rst content with proper title formatting + # Title should use dots as separators, but escape underscores for RST + title = module_name.replace('_', r'\_') + underline = '=' * len(title) # Underline matches title length + + rst_content = f"""{title} +{underline} + +.. automodule:: {module_name} + :members: + :undoc-members: + :show-inheritance: + :member-order: bysource +""" + + # Write the .rst file + with open(rst_path, 'w') as f: + f.write(rst_content) + + generated_files.add(module_name) + + return generated_files + + +def generate_toctree_from_rst_files(reference_dir: Path) -> Dict[str, List[str]]: + """Generate toctree entries directly from existing .rst files.""" + # Initialize sections using constants + toctree_sections: Dict[str, List[str]] = {section: [] for section in PACKAGE_SECTIONS.values()} + + python_ref_dir = reference_dir / "python" + if not python_ref_dir.exists(): + return toctree_sections + + # Collect modules by package using constants + modules_by_section: Dict[str, List[str]] = {section: [] for section in PACKAGE_SECTIONS.values()} + + # Get all .rst files and organize them by package + for rst_file in python_ref_dir.glob("*.rst"): + module_name = rst_file.stem # filename without .rst extension + + # Find which documented package this module belongs to + for package_prefix, section_name in PACKAGE_SECTIONS.items(): + if module_name.startswith(package_prefix): + modules_by_section[section_name].append(module_name) + break + + # Sort modules so parent modules come before child modules + def sort_modules_hierarchically(modules): + """Sort modules so that parent modules come before child modules.""" + return sorted(modules, key=lambda x: (x.count('.'), x)) + + # Apply hierarchical sorting and convert to rst paths + for section_name, modules in modules_by_section.items(): + toctree_sections[section_name] = [f"python/{m}" for m in sort_modules_hierarchically(modules)] + + return toctree_sections + + +def generate_index_content(toctree_sections: Dict[str, List[str]]) -> str: + """Generate the complete index.md content with automatic toctrees.""" + + content = """--- +myst: + html_meta: + "description lang=en": | + AutoGen is a community-driven project. Learn how to get involved, contribute, and connect with the community. +--- + +# API Reference + +""" + + for section_name, modules in toctree_sections.items(): + if modules: # Only add section if it has modules + content += f"""```{{toctree}} +:caption: {section_name} +:maxdepth: 2 + +""" + for module in modules: + content += f"{module}\n" + content += "```\n\n" + + return content + + +def main(): + """Main function to generate the API documentation index.""" + script_dir = Path(__file__).parent + reference_dir = script_dir / "reference" + index_file = reference_dir / "index.md" + + print("🔍 Scanning Python packages...") + package_roots = find_python_packages() + + all_modules = {} + for package_root in package_roots: + print(f" 📦 Scanning {package_root}") + modules = get_module_hierarchy(package_root) + all_modules.update(modules) + + print("🏗️ Generating .rst files for all discovered modules...") + generated_files = generate_rst_files(package_roots, reference_dir) + print(f" Generated {len(generated_files)} .rst files") + + print("📝 Generating toctree entries from .rst files...") + toctree_sections = generate_toctree_from_rst_files(reference_dir) + + for section, modules in toctree_sections.items(): + print(f" {section}: {len(modules)} modules") + + print("✍️ Writing index.md...") + content = generate_index_content(toctree_sections) + + with open(index_file, 'w') as f: + f.write(content) + + print(f"✅ Generated API documentation index at {index_file}") + print("\n📖 Summary:") + total_modules = sum(len(modules) for modules in toctree_sections.values()) + print(f" Total modules documented: {total_modules}") + + for section, modules in toctree_sections.items(): + if modules: + print(f" {section}: {len(modules)} modules") + + +if __name__ == "__main__": + main() diff --git a/python/docs/src/images/assistant-agent.svg b/python/docs/src/images/assistant-agent.svg new file mode 100644 index 000000000000..dd16f69020ab --- /dev/null +++ b/python/docs/src/images/assistant-agent.svg @@ -0,0 +1,3 @@ + + +
Model Context
Model Context
Model Context
Model Context
New Messages
New Messages
Memory
Memory
Model Client
Model Client
Tools
Tools
Model Context
Model Context
1. Add New Messages to Context
1. Add New Messages to Context
Model Context
Model Context
2. Update Context
2. Update Context
Model Context
Model Context
3. Chat Completion
3. Chat Completion
Model Context
Model Context
Model Context
Model Context
Model Context
Model Context
4. Tool Execution
4. Tool Execution
Model Client
Model Client
5. Chat Completion (Reflect on Tool Use)
5. Chat Completion (R...
Model Context
Model Context
Model Context
Model Context
Model Context
Model Context
Model Context
Model Context
Response
(Text)
Response...
Response 
(Tool Result Summary)
Response...
Response 
(Text Message)
Response...
Model Context
Model Context
Model Context
Model Context
Model Context
Model Context
Tool Call Detected?
Tool Call Detected?
Reflect on Tool Use?
Reflect on Tool Use?
No
No
No
No
Handoff Detected?
Handoff Detected?
Response
(Handoff)
Response...
Yes
Yes
Assistant Agent
Assistant Agent
Maximum Tool Iterations Reached?
Maximum Tool Iterations Reached?
Yes
Yes
No
No \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/images/autogen-magentic-one-agents.png b/python/docs/src/images/autogen-magentic-one-agents.png similarity index 100% rename from python/packages/autogen-core/docs/src/images/autogen-magentic-one-agents.png rename to python/docs/src/images/autogen-magentic-one-agents.png diff --git a/python/packages/autogen-core/docs/src/images/autogen-magentic-one-example.png b/python/docs/src/images/autogen-magentic-one-example.png similarity index 100% rename from python/packages/autogen-core/docs/src/images/autogen-magentic-one-example.png rename to python/docs/src/images/autogen-magentic-one-example.png diff --git a/python/packages/autogen-core/docs/src/images/code.svg b/python/docs/src/images/code.svg similarity index 100% rename from python/packages/autogen-core/docs/src/images/code.svg rename to python/docs/src/images/code.svg diff --git a/python/packages/autogen-core/docs/src/images/example-company.jpg b/python/docs/src/images/example-company.jpg similarity index 100% rename from python/packages/autogen-core/docs/src/images/example-company.jpg rename to python/docs/src/images/example-company.jpg diff --git a/python/packages/autogen-core/docs/src/images/example-literature.jpg b/python/docs/src/images/example-literature.jpg similarity index 100% rename from python/packages/autogen-core/docs/src/images/example-literature.jpg rename to python/docs/src/images/example-literature.jpg diff --git a/python/packages/autogen-core/docs/src/images/example-travel.jpeg b/python/docs/src/images/example-travel.jpeg similarity index 100% rename from python/packages/autogen-core/docs/src/images/example-travel.jpeg rename to python/docs/src/images/example-travel.jpeg diff --git a/python/packages/autogen-core/docs/src/images/open-ai-telemetry-example.png b/python/docs/src/images/open-ai-telemetry-example.png similarity index 100% rename from python/packages/autogen-core/docs/src/images/open-ai-telemetry-example.png rename to python/docs/src/images/open-ai-telemetry-example.png diff --git a/python/packages/autogen-core/docs/src/index.md b/python/docs/src/index.md similarity index 80% rename from python/packages/autogen-core/docs/src/index.md rename to python/docs/src/index.md index 1f13cc8dc237..ca41ee79c32f 100644 --- a/python/packages/autogen-core/docs/src/index.md +++ b/python/docs/src/index.md @@ -46,40 +46,12 @@ A framework for building AI agents and applications ::::{grid} :gutter: 2 -:::{grid-item-card} -:shadow: none -:margin: 2 0 0 0 -:columns: 12 12 6 6 - -
- -{fas}`book;pst-color-primary` -Magentic-One CLI [![PyPi magentic-one-cli](https://img.shields.io/badge/PyPi-magentic--one--cli-blue?logo=pypi)](https://pypi.org/project/magentic-one-cli/) -
-A console-based multi-agent assistant for web and file-based tasks. -Built on AgentChat. - -```bash -pip install -U magentic-one-cli -m1 "Find flights from Seattle to Paris and format the result in a table" -``` - -+++ - -```{button-ref} user-guide/agentchat-user-guide/magentic-one -:color: secondary - -Get Started -``` - -::: - :::{grid-item-card} {fas}`palette;pst-color-primary` Studio [![PyPi autogenstudio](https://img.shields.io/badge/PyPi-autogenstudio-blue?logo=pypi)](https://pypi.org/project/autogenstudio/) :shadow: none :margin: 2 0 0 0 -:columns: 12 12 6 6 +:columns: 12 12 12 12 -An app for prototyping and managing agents without writing code. +An web-based UI for prototyping with agents without writing code. Built on AgentChat. ```bash @@ -87,6 +59,8 @@ pip install -U autogenstudio autogenstudio ui --port 8080 --appdir ./myapp ``` +_Start here if you are new to AutoGen and want to prototype with agents without writing code._ + +++ ```{button-ref} user-guide/autogenstudio-user-guide/index @@ -124,7 +98,7 @@ async def main() -> None: asyncio.run(main()) ``` -_Start here if you are building conversational agents. [Migrating from AutoGen 0.2?](./user-guide/agentchat-user-guide/migration-guide.md)._ +_Start here if you are prototyping with agents using Python. [Migrating from AutoGen 0.2?](./user-guide/agentchat-user-guide/migration-guide.md)._ +++ @@ -147,7 +121,7 @@ An event-driven programming framework for building scalable multi-agent AI syste * Research on multi-agent collaboration. * Distributed agents for multi-language applications. -_Start here if you are building workflows or distributed agent systems._ +_Start here if you are getting serious about building multi-agent systems._ +++ @@ -167,7 +141,7 @@ Get Started Implementations of Core and AgentChat components that interface with external services or other libraries. You can find and use community extensions or create your own. Examples of built-in extensions: -* {py:class}`~autogen_ext.tools.langchain.LangChainToolAdapter` for using LangChain tools. +* {py:class}`~autogen_ext.tools.mcp.McpWorkbench` for using Model-Context Protocol (MCP) servers. * {py:class}`~autogen_ext.agents.openai.OpenAIAssistantAgent` for using Assistant API. * {py:class}`~autogen_ext.code_executors.docker.DockerCommandLineCodeExecutor` for running model-generated code in a Docker container. * {py:class}`~autogen_ext.runtimes.grpc.GrpcWorkerAgentRuntime` for distributed agents. diff --git a/python/docs/src/user-guide/agentchat-user-guide/custom-agents.ipynb b/python/docs/src/user-guide/agentchat-user-guide/custom-agents.ipynb new file mode 100644 index 000000000000..b24b49e76cd2 --- /dev/null +++ b/python/docs/src/user-guide/agentchat-user-guide/custom-agents.ipynb @@ -0,0 +1,741 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Custom Agents\n", + "\n", + "You may have agents with behaviors that do not fall into a preset. \n", + "In such cases, you can build custom agents.\n", + "\n", + "All agents in AgentChat inherit from {py:class}`~autogen_agentchat.agents.BaseChatAgent` \n", + "class and implement the following abstract methods and attributes:\n", + "\n", + "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages`: The abstract method that defines the behavior of the agent in response to messages. This method is called when the agent is asked to provide a response in {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run`. It returns a {py:class}`~autogen_agentchat.base.Response` object.\n", + "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_reset`: The abstract method that resets the agent to its initial state. This method is called when the agent is asked to reset itself.\n", + "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.produced_message_types`: The list of possible {py:class}`~autogen_agentchat.messages.BaseChatMessage` message types the agent can produce in its response.\n", + "\n", + "Optionally, you can implement the the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream` method to stream messages as they are generated by the agent.\n", + "This method is called by {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run_stream` to stream messages.\n", + "If this method is not implemented, the agent\n", + "uses the default implementation of {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream`\n", + "that calls the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages` method and\n", + "yields all messages in the response." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## CountDownAgent\n", + "\n", + "In this example, we create a simple agent that counts down from a given number to zero,\n", + "and produces a stream of messages with the current count." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3...\n", + "2...\n", + "1...\n", + "Done!\n" + ] + } + ], + "source": [ + "from typing import AsyncGenerator, List, Sequence\n", + "\n", + "from autogen_agentchat.agents import BaseChatAgent\n", + "from autogen_agentchat.base import Response\n", + "from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage, TextMessage\n", + "from autogen_core import CancellationToken\n", + "\n", + "\n", + "class CountDownAgent(BaseChatAgent):\n", + " def __init__(self, name: str, count: int = 3):\n", + " super().__init__(name, \"A simple agent that counts down.\")\n", + " self._count = count\n", + "\n", + " @property\n", + " def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:\n", + " return (TextMessage,)\n", + "\n", + " async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:\n", + " # Calls the on_messages_stream.\n", + " response: Response | None = None\n", + " async for message in self.on_messages_stream(messages, cancellation_token):\n", + " if isinstance(message, Response):\n", + " response = message\n", + " assert response is not None\n", + " return response\n", + "\n", + " async def on_messages_stream(\n", + " self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken\n", + " ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:\n", + " inner_messages: List[BaseAgentEvent | BaseChatMessage] = []\n", + " for i in range(self._count, 0, -1):\n", + " msg = TextMessage(content=f\"{i}...\", source=self.name)\n", + " inner_messages.append(msg)\n", + " yield msg\n", + " # The response is returned at the end of the stream.\n", + " # It contains the final message and all the inner messages.\n", + " yield Response(chat_message=TextMessage(content=\"Done!\", source=self.name), inner_messages=inner_messages)\n", + "\n", + " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", + " pass\n", + "\n", + "\n", + "async def run_countdown_agent() -> None:\n", + " # Create a countdown agent.\n", + " countdown_agent = CountDownAgent(\"countdown\")\n", + "\n", + " # Run the agent with a given task and stream the response.\n", + " async for message in countdown_agent.on_messages_stream([], CancellationToken()):\n", + " if isinstance(message, Response):\n", + " print(message.chat_message)\n", + " else:\n", + " print(message)\n", + "\n", + "\n", + "# Use asyncio.run(run_countdown_agent()) when running in a script.\n", + "await run_countdown_agent()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## ArithmeticAgent\n", + "\n", + "In this example, we create an agent class that can perform simple arithmetic operations\n", + "on a given integer. Then, we will use different instances of this agent class\n", + "in a {py:class}`~autogen_agentchat.teams.SelectorGroupChat`\n", + "to transform a given integer into another integer by applying a sequence of arithmetic operations.\n", + "\n", + "The `ArithmeticAgent` class takes an `operator_func` that takes an integer and returns an integer,\n", + "after applying an arithmetic operation to the integer.\n", + "In its `on_messages` method, it applies the `operator_func` to the integer in the input message,\n", + "and returns a response with the result." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Callable, Sequence\n", + "\n", + "from autogen_agentchat.agents import BaseChatAgent\n", + "from autogen_agentchat.base import Response\n", + "from autogen_agentchat.conditions import MaxMessageTermination\n", + "from autogen_agentchat.messages import BaseChatMessage\n", + "from autogen_agentchat.teams import SelectorGroupChat\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_core import CancellationToken\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "\n", + "class ArithmeticAgent(BaseChatAgent):\n", + " def __init__(self, name: str, description: str, operator_func: Callable[[int], int]) -> None:\n", + " super().__init__(name, description=description)\n", + " self._operator_func = operator_func\n", + " self._message_history: List[BaseChatMessage] = []\n", + "\n", + " @property\n", + " def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:\n", + " return (TextMessage,)\n", + "\n", + " async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:\n", + " # Update the message history.\n", + " # NOTE: it is possible the messages is an empty list, which means the agent was selected previously.\n", + " self._message_history.extend(messages)\n", + " # Parse the number in the last message.\n", + " assert isinstance(self._message_history[-1], TextMessage)\n", + " number = int(self._message_history[-1].content)\n", + " # Apply the operator function to the number.\n", + " result = self._operator_func(number)\n", + " # Create a new message with the result.\n", + " response_message = TextMessage(content=str(result), source=self.name)\n", + " # Update the message history.\n", + " self._message_history.append(response_message)\n", + " # Return the response.\n", + " return Response(chat_message=response_message)\n", + "\n", + " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", + " pass" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "```{note}\n", + "The `on_messages` method may be called with an empty list of messages, in which\n", + "case it means the agent was called previously and is now being called again,\n", + "without any new messages from the caller. So it is important to keep a history\n", + "of the previous messages received by the agent, and use that history to generate\n", + "the response.\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we can create a {py:class}`~autogen_agentchat.teams.SelectorGroupChat` with 5 instances of `ArithmeticAgent`:\n", + "\n", + "- one that adds 1 to the input integer,\n", + "- one that subtracts 1 from the input integer,\n", + "- one that multiplies the input integer by 2,\n", + "- one that divides the input integer by 2 and rounds down to the nearest integer, and\n", + "- one that returns the input integer unchanged.\n", + "\n", + "We then create a {py:class}`~autogen_agentchat.teams.SelectorGroupChat` with these agents,\n", + "and set the appropriate selector settings:\n", + "\n", + "- allow the same agent to be selected consecutively to allow for repeated operations, and\n", + "- customize the selector prompt to tailor the model's response to the specific task." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "Apply the operations to turn the given number into 25.\n", + "---------- user ----------\n", + "10\n", + "---------- multiply_agent ----------\n", + "20\n", + "---------- add_agent ----------\n", + "21\n", + "---------- multiply_agent ----------\n", + "42\n", + "---------- divide_agent ----------\n", + "21\n", + "---------- add_agent ----------\n", + "22\n", + "---------- add_agent ----------\n", + "23\n", + "---------- add_agent ----------\n", + "24\n", + "---------- add_agent ----------\n", + "25\n", + "---------- Summary ----------\n", + "Number of messages: 10\n", + "Finish reason: Maximum number of messages 10 reached, current message count: 10\n", + "Total prompt tokens: 0\n", + "Total completion tokens: 0\n", + "Duration: 2.40 seconds\n" + ] + } + ], + "source": [ + "async def run_number_agents() -> None:\n", + " # Create agents for number operations.\n", + " add_agent = ArithmeticAgent(\"add_agent\", \"Adds 1 to the number.\", lambda x: x + 1)\n", + " multiply_agent = ArithmeticAgent(\"multiply_agent\", \"Multiplies the number by 2.\", lambda x: x * 2)\n", + " subtract_agent = ArithmeticAgent(\"subtract_agent\", \"Subtracts 1 from the number.\", lambda x: x - 1)\n", + " divide_agent = ArithmeticAgent(\"divide_agent\", \"Divides the number by 2 and rounds down.\", lambda x: x // 2)\n", + " identity_agent = ArithmeticAgent(\"identity_agent\", \"Returns the number as is.\", lambda x: x)\n", + "\n", + " # The termination condition is to stop after 10 messages.\n", + " termination_condition = MaxMessageTermination(10)\n", + "\n", + " # Create a selector group chat.\n", + " selector_group_chat = SelectorGroupChat(\n", + " [add_agent, multiply_agent, subtract_agent, divide_agent, identity_agent],\n", + " model_client=OpenAIChatCompletionClient(model=\"gpt-4o\"),\n", + " termination_condition=termination_condition,\n", + " allow_repeated_speaker=True, # Allow the same agent to speak multiple times, necessary for this task.\n", + " selector_prompt=(\n", + " \"Available roles:\\n{roles}\\nTheir job descriptions:\\n{participants}\\n\"\n", + " \"Current conversation history:\\n{history}\\n\"\n", + " \"Please select the most appropriate role for the next message, and only return the role name.\"\n", + " ),\n", + " )\n", + "\n", + " # Run the selector group chat with a given task and stream the response.\n", + " task: List[BaseChatMessage] = [\n", + " TextMessage(content=\"Apply the operations to turn the given number into 25.\", source=\"user\"),\n", + " TextMessage(content=\"10\", source=\"user\"),\n", + " ]\n", + " stream = selector_group_chat.run_stream(task=task)\n", + " await Console(stream)\n", + "\n", + "\n", + "# Use asyncio.run(run_number_agents()) when running in a script.\n", + "await run_number_agents()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From the output, we can see that the agents have successfully transformed the input integer\n", + "from 10 to 25 by choosing appropriate agents that apply the arithmetic operations in sequence." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Custom Model Clients in Custom Agents\n", + "\n", + "One of the key features of the {py:class}`~autogen_agentchat.agents.AssistantAgent` preset in AgentChat is that it takes a `model_client` argument and can use it in responding to messages. However, in some cases, you may want your agent to use a custom model client that is not currently supported (see [supported model clients](https://microsoft.github.io/autogen/dev/user-guide/core-user-guide/components/model-clients.html)) or custom model behaviours. \n", + "\n", + "You can accomplish this with a custom agent that implements *your custom model client*.\n", + "\n", + "In the example below, we will walk through an example of a custom agent that uses the [Google Gemini SDK](https://github.com/googleapis/python-genai) directly to respond to messages.\n", + "\n", + "> **Note:** You will need to install the [Google Gemini SDK](https://github.com/googleapis/python-genai) to run this example. You can install it using the following command: \n", + "\n", + "```bash\n", + "pip install google-genai\n", + "``` " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# !pip install google-genai\n", + "import os\n", + "from typing import AsyncGenerator, Sequence\n", + "\n", + "from autogen_agentchat.agents import BaseChatAgent\n", + "from autogen_agentchat.base import Response\n", + "from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage\n", + "from autogen_core import CancellationToken\n", + "from autogen_core.model_context import UnboundedChatCompletionContext\n", + "from autogen_core.models import AssistantMessage, RequestUsage, UserMessage\n", + "from google import genai\n", + "from google.genai import types\n", + "\n", + "\n", + "class GeminiAssistantAgent(BaseChatAgent):\n", + " def __init__(\n", + " self,\n", + " name: str,\n", + " description: str = \"An agent that provides assistance with ability to use tools.\",\n", + " model: str = \"gemini-1.5-flash-002\",\n", + " api_key: str = os.environ[\"GEMINI_API_KEY\"],\n", + " system_message: str\n", + " | None = \"You are a helpful assistant that can respond to messages. Reply with TERMINATE when the task has been completed.\",\n", + " ):\n", + " super().__init__(name=name, description=description)\n", + " self._model_context = UnboundedChatCompletionContext()\n", + " self._model_client = genai.Client(api_key=api_key)\n", + " self._system_message = system_message\n", + " self._model = model\n", + "\n", + " @property\n", + " def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:\n", + " return (TextMessage,)\n", + "\n", + " async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:\n", + " final_response = None\n", + " async for message in self.on_messages_stream(messages, cancellation_token):\n", + " if isinstance(message, Response):\n", + " final_response = message\n", + "\n", + " if final_response is None:\n", + " raise AssertionError(\"The stream should have returned the final result.\")\n", + "\n", + " return final_response\n", + "\n", + " async def on_messages_stream(\n", + " self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken\n", + " ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:\n", + " # Add messages to the model context\n", + " for msg in messages:\n", + " await self._model_context.add_message(msg.to_model_message())\n", + "\n", + " # Get conversation history\n", + " history = [\n", + " (msg.source if hasattr(msg, \"source\") else \"system\")\n", + " + \": \"\n", + " + (msg.content if isinstance(msg.content, str) else \"\")\n", + " + \"\\n\"\n", + " for msg in await self._model_context.get_messages()\n", + " ]\n", + " # Generate response using Gemini\n", + " response = self._model_client.models.generate_content(\n", + " model=self._model,\n", + " contents=f\"History: {history}\\nGiven the history, please provide a response\",\n", + " config=types.GenerateContentConfig(\n", + " system_instruction=self._system_message,\n", + " temperature=0.3,\n", + " ),\n", + " )\n", + "\n", + " # Create usage metadata\n", + " usage = RequestUsage(\n", + " prompt_tokens=response.usage_metadata.prompt_token_count,\n", + " completion_tokens=response.usage_metadata.candidates_token_count,\n", + " )\n", + "\n", + " # Add response to model context\n", + " await self._model_context.add_message(AssistantMessage(content=response.text, source=self.name))\n", + "\n", + " # Yield the final response\n", + " yield Response(\n", + " chat_message=TextMessage(content=response.text, source=self.name, models_usage=usage),\n", + " inner_messages=[],\n", + " )\n", + "\n", + " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", + " \"\"\"Reset the assistant by clearing the model context.\"\"\"\n", + " await self._model_context.clear()" + ] + }, + { + "cell_type": "code", + "execution_count": 38, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "What is the capital of New York?\n", + "---------- gemini_assistant ----------\n", + "Albany\n", + "TERMINATE\n", + "\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What is the capital of New York?', type='TextMessage'), TextMessage(source='gemini_assistant', models_usage=RequestUsage(prompt_tokens=46, completion_tokens=5), content='Albany\\nTERMINATE\\n', type='TextMessage')], stop_reason=None)" + ] + }, + "execution_count": 38, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "gemini_assistant = GeminiAssistantAgent(\"gemini_assistant\")\n", + "await Console(gemini_assistant.run_stream(task=\"What is the capital of New York?\"))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the example above, we have chosen to provide `model`, `api_key` and `system_message` as arguments - you can choose to provide any other arguments that are required by the model client you are using or fits with your application design. \n", + "\n", + "Now, let us explore how to use this custom agent as part of a team in AgentChat." + ] + }, + { + "cell_type": "code", + "execution_count": 39, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "Write a Haiku poem with 4 lines about the fall season.\n", + "---------- primary ----------\n", + "Crimson leaves cascade, \n", + "Whispering winds sing of change, \n", + "Chill wraps the fading, \n", + "Nature's quilt, rich and warm.\n", + "---------- gemini_critic ----------\n", + "The poem is good, but it has four lines instead of three. A haiku must have three lines with a 5-7-5 syllable structure. The content is evocative of autumn, but the form is incorrect. Please revise to adhere to the haiku's syllable structure.\n", + "\n", + "---------- primary ----------\n", + "Thank you for your feedback! Here’s a revised haiku that follows the 5-7-5 syllable structure:\n", + "\n", + "Crimson leaves drift down, \n", + "Chill winds whisper through the gold, \n", + "Autumn’s breath is near.\n", + "---------- gemini_critic ----------\n", + "The revised haiku is much improved. It correctly follows the 5-7-5 syllable structure and maintains the evocative imagery of autumn. APPROVE\n", + "\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Write a Haiku poem with 4 lines about the fall season.', type='TextMessage'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=33, completion_tokens=31), content=\"Crimson leaves cascade, \\nWhispering winds sing of change, \\nChill wraps the fading, \\nNature's quilt, rich and warm.\", type='TextMessage'), TextMessage(source='gemini_critic', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=60), content=\"The poem is good, but it has four lines instead of three. A haiku must have three lines with a 5-7-5 syllable structure. The content is evocative of autumn, but the form is incorrect. Please revise to adhere to the haiku's syllable structure.\\n\", type='TextMessage'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=141, completion_tokens=49), content='Thank you for your feedback! Here’s a revised haiku that follows the 5-7-5 syllable structure:\\n\\nCrimson leaves drift down, \\nChill winds whisper through the gold, \\nAutumn’s breath is near.', type='TextMessage'), TextMessage(source='gemini_critic', models_usage=RequestUsage(prompt_tokens=211, completion_tokens=32), content='The revised haiku is much improved. It correctly follows the 5-7-5 syllable structure and maintains the evocative imagery of autumn. APPROVE\\n', type='TextMessage')], stop_reason=\"Text 'APPROVE' mentioned\")" + ] + }, + "execution_count": 39, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.conditions import TextMentionTermination\n", + "from autogen_agentchat.teams import RoundRobinGroupChat\n", + "from autogen_agentchat.ui import Console\n", + "\n", + "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\")\n", + "\n", + "# Create the primary agent.\n", + "primary_agent = AssistantAgent(\n", + " \"primary\",\n", + " model_client=model_client,\n", + " system_message=\"You are a helpful AI assistant.\",\n", + ")\n", + "\n", + "# Create a critic agent based on our new GeminiAssistantAgent.\n", + "gemini_critic_agent = GeminiAssistantAgent(\n", + " \"gemini_critic\",\n", + " system_message=\"Provide constructive feedback. Respond with 'APPROVE' to when your feedbacks are addressed.\",\n", + ")\n", + "\n", + "\n", + "# Define a termination condition that stops the task if the critic approves or after 10 messages.\n", + "termination = TextMentionTermination(\"APPROVE\") | MaxMessageTermination(10)\n", + "\n", + "# Create a team with the primary and critic agents.\n", + "team = RoundRobinGroupChat([primary_agent, gemini_critic_agent], termination_condition=termination)\n", + "\n", + "await Console(team.run_stream(task=\"Write a Haiku poem with 4 lines about the fall season.\"))\n", + "await model_client.close()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In section above, we show several very important concepts:\n", + "- We have developed a custom agent that uses the Google Gemini SDK to respond to messages. \n", + "- We show that this custom agent can be used as part of the broader AgentChat ecosystem - in this case as a participant in a {py:class}`~autogen_agentchat.teams.RoundRobinGroupChat` as long as it inherits from {py:class}`~autogen_agentchat.agents.BaseChatAgent`.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Making the Custom Agent Declarative \n", + "\n", + "Autogen provides a [Component](https://microsoft.github.io/autogen/dev/user-guide/core-user-guide/framework/component-config.html) interface for making the configuration of components serializable to a declarative format. This is useful for saving and loading configurations, and for sharing configurations with others. \n", + "\n", + "We accomplish this by inheriting from the `Component` class and implementing the `_from_config` and `_to_config` methods.\n", + "The declarative class can be serialized to a JSON format using the `dump_component` method, and deserialized from a JSON format using the `load_component` method." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from typing import AsyncGenerator, Sequence\n", + "\n", + "from autogen_agentchat.agents import BaseChatAgent\n", + "from autogen_agentchat.base import Response\n", + "from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage\n", + "from autogen_core import CancellationToken, Component\n", + "from pydantic import BaseModel\n", + "from typing_extensions import Self\n", + "\n", + "\n", + "class GeminiAssistantAgentConfig(BaseModel):\n", + " name: str\n", + " description: str = \"An agent that provides assistance with ability to use tools.\"\n", + " model: str = \"gemini-1.5-flash-002\"\n", + " system_message: str | None = None\n", + "\n", + "\n", + "class GeminiAssistantAgent(BaseChatAgent, Component[GeminiAssistantAgentConfig]): # type: ignore[no-redef]\n", + " component_config_schema = GeminiAssistantAgentConfig\n", + " # component_provider_override = \"mypackage.agents.GeminiAssistantAgent\"\n", + "\n", + " def __init__(\n", + " self,\n", + " name: str,\n", + " description: str = \"An agent that provides assistance with ability to use tools.\",\n", + " model: str = \"gemini-1.5-flash-002\",\n", + " api_key: str = os.environ[\"GEMINI_API_KEY\"],\n", + " system_message: str\n", + " | None = \"You are a helpful assistant that can respond to messages. Reply with TERMINATE when the task has been completed.\",\n", + " ):\n", + " super().__init__(name=name, description=description)\n", + " self._model_context = UnboundedChatCompletionContext()\n", + " self._model_client = genai.Client(api_key=api_key)\n", + " self._system_message = system_message\n", + " self._model = model\n", + "\n", + " @property\n", + " def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:\n", + " return (TextMessage,)\n", + "\n", + " async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:\n", + " final_response = None\n", + " async for message in self.on_messages_stream(messages, cancellation_token):\n", + " if isinstance(message, Response):\n", + " final_response = message\n", + "\n", + " if final_response is None:\n", + " raise AssertionError(\"The stream should have returned the final result.\")\n", + "\n", + " return final_response\n", + "\n", + " async def on_messages_stream(\n", + " self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken\n", + " ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:\n", + " # Add messages to the model context\n", + " for msg in messages:\n", + " await self._model_context.add_message(msg.to_model_message())\n", + "\n", + " # Get conversation history\n", + " history = [\n", + " (msg.source if hasattr(msg, \"source\") else \"system\")\n", + " + \": \"\n", + " + (msg.content if isinstance(msg.content, str) else \"\")\n", + " + \"\\n\"\n", + " for msg in await self._model_context.get_messages()\n", + " ]\n", + "\n", + " # Generate response using Gemini\n", + " response = self._model_client.models.generate_content(\n", + " model=self._model,\n", + " contents=f\"History: {history}\\nGiven the history, please provide a response\",\n", + " config=types.GenerateContentConfig(\n", + " system_instruction=self._system_message,\n", + " temperature=0.3,\n", + " ),\n", + " )\n", + "\n", + " # Create usage metadata\n", + " usage = RequestUsage(\n", + " prompt_tokens=response.usage_metadata.prompt_token_count,\n", + " completion_tokens=response.usage_metadata.candidates_token_count,\n", + " )\n", + "\n", + " # Add response to model context\n", + " await self._model_context.add_message(AssistantMessage(content=response.text, source=self.name))\n", + "\n", + " # Yield the final response\n", + " yield Response(\n", + " chat_message=TextMessage(content=response.text, source=self.name, models_usage=usage),\n", + " inner_messages=[],\n", + " )\n", + "\n", + " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", + " \"\"\"Reset the assistant by clearing the model context.\"\"\"\n", + " await self._model_context.clear()\n", + "\n", + " @classmethod\n", + " def _from_config(cls, config: GeminiAssistantAgentConfig) -> Self:\n", + " return cls(\n", + " name=config.name, description=config.description, model=config.model, system_message=config.system_message\n", + " )\n", + "\n", + " def _to_config(self) -> GeminiAssistantAgentConfig:\n", + " return GeminiAssistantAgentConfig(\n", + " name=self.name,\n", + " description=self.description,\n", + " model=self._model,\n", + " system_message=self._system_message,\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now that we have the required methods implemented, we can now load and dump the custom agent to and from a JSON format, and then load the agent from the JSON format.\n", + " \n", + " > Note: You should set the `component_provider_override` class variable to the full path of the module containing the custom agent class e.g., (`mypackage.agents.GeminiAssistantAgent`). This is used by `load_component` method to determine how to instantiate the class. \n", + " " + ] + }, + { + "cell_type": "code", + "execution_count": 41, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{\n", + " \"provider\": \"__main__.GeminiAssistantAgent\",\n", + " \"component_type\": \"agent\",\n", + " \"version\": 1,\n", + " \"component_version\": 1,\n", + " \"description\": null,\n", + " \"label\": \"GeminiAssistantAgent\",\n", + " \"config\": {\n", + " \"name\": \"gemini_assistant\",\n", + " \"description\": \"An agent that provides assistance with ability to use tools.\",\n", + " \"model\": \"gemini-1.5-flash-002\",\n", + " \"system_message\": \"You are a helpful assistant that can respond to messages. Reply with TERMINATE when the task has been completed.\"\n", + " }\n", + "}\n", + "<__main__.GeminiAssistantAgent object at 0x11a5c5a90>\n" + ] + } + ], + "source": [ + "gemini_assistant = GeminiAssistantAgent(\"gemini_assistant\")\n", + "config = gemini_assistant.dump_component()\n", + "print(config.model_dump_json(indent=2))\n", + "loaded_agent = GeminiAssistantAgent.load_component(config)\n", + "print(loaded_agent)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Next Steps \n", + "\n", + "So far, we have seen how to create custom agents, add custom model clients to agents, and make custom agents declarative. There are a few ways in which this basic sample can be extended:\n", + "\n", + "- Extend the Gemini model client to handle function calling similar to the {py:class}`~autogen_agentchat.agents.AssistantAgent` class. https://ai.google.dev/gemini-api/docs/function-calling \n", + "- Implement a package with a custom agent and experiment with using its declarative format in a tool like [AutoGen Studio](https://microsoft.github.io/autogen/stable/user-guide/autogenstudio-user-guide/index.html)." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.7" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/examples/company-research.ipynb b/python/docs/src/user-guide/agentchat-user-guide/examples/company-research.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/examples/company-research.ipynb rename to python/docs/src/user-guide/agentchat-user-guide/examples/company-research.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/examples/index.md b/python/docs/src/user-guide/agentchat-user-guide/examples/index.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/examples/index.md rename to python/docs/src/user-guide/agentchat-user-guide/examples/index.md diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/examples/literature-review.ipynb b/python/docs/src/user-guide/agentchat-user-guide/examples/literature-review.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/examples/literature-review.ipynb rename to python/docs/src/user-guide/agentchat-user-guide/examples/literature-review.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/examples/travel-planning.ipynb b/python/docs/src/user-guide/agentchat-user-guide/examples/travel-planning.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/examples/travel-planning.ipynb rename to python/docs/src/user-guide/agentchat-user-guide/examples/travel-planning.ipynb diff --git a/python/docs/src/user-guide/agentchat-user-guide/graph-flow.ipynb b/python/docs/src/user-guide/agentchat-user-guide/graph-flow.ipynb new file mode 100644 index 000000000000..34063ed78582 --- /dev/null +++ b/python/docs/src/user-guide/agentchat-user-guide/graph-flow.ipynb @@ -0,0 +1,825 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "b3b5a600", + "metadata": {}, + "source": [ + "# GraphFlow (Workflows)\n", + "\n", + "In this section you'll learn how to create an _multi-agent workflow_ using {py:class}`~autogen_agentchat.teams.GraphFlow`, or simply \"flow\" for short.\n", + "It uses structured execution and precisely controls how agents interact to accomplish a task.\n", + "\n", + "We'll first show you how to create and run a flow. We'll then explain how to observe and debug flow behavior, \n", + "and discuss important operations for managing execution.\n", + "\n", + "AutoGen AgentChat provides a team for directed graph execution:\n", + "\n", + "- {py:class}`~autogen_agentchat.teams.GraphFlow`: A team that follows a {py:class}`~autogen_agentchat.teams.DiGraph`\n", + "to control the execution flow between agents. \n", + "Supports sequential, parallel, conditional, and looping behaviors.\n", + "\n", + "```{note}\n", + "**When should you use {py:class}`~autogen_agentchat.teams.GraphFlow`?**\n", + "\n", + "Use Graph when you need strict control over the order in which agents act, or when different outcomes must lead to different next steps.\n", + "Start with a simple team such as {py:class}`~autogen_agentchat.teams.RoundRobinGroupChat` or {py:class}`~autogen_agentchat.teams.SelectorGroupChat`\n", + "if ad-hoc conversation flow is sufficient. \n", + "Transition to a structured workflow when your task requires deterministic control,\n", + "conditional branching, or handling complex multi-step processes with cycles.\n", + "```\n", + "\n", + "> **Warning:** {py:class}`~autogen_agentchat.teams.GraphFlow` is an **experimental feature**. \n", + "Its API, behavior, and capabilities are **subject to change** in future releases." + ] + }, + { + "cell_type": "markdown", + "id": "f04e4271", + "metadata": {}, + "source": [ + "## Creating and Running a Flow\n", + "\n", + "{py:class}`~autogen_agentchat.teams.DiGraphBuilder` is a fluent utility that lets you easily construct execution graphs for workflows. It supports building:\n", + "\n", + "- Sequential chains\n", + "- Parallel fan-outs\n", + "- Conditional branching\n", + "- Loops with safe exit conditions\n", + "\n", + "Each node in the graph represents an agent, and edges define the allowed execution paths. Edges can optionally have conditions based on agent messages." + ] + }, + { + "cell_type": "markdown", + "id": "760ee783", + "metadata": {}, + "source": [ + "### Sequential Flow\n", + "\n", + "We will begin by creating a simple workflow where a **writer** drafts a paragraph and a **reviewer** provides feedback. This graph terminates after the reviewer comments on the writer. \n", + "\n", + "Note, the flow automatically computes all the source and leaf nodes of the graph and the execution starts at all the source nodes in the graph and completes execution when no nodes are left to execute." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "c6dd0386", + "metadata": {}, + "outputs": [], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.teams import DiGraphBuilder, GraphFlow\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "# Create an OpenAI model client\n", + "client = OpenAIChatCompletionClient(model=\"gpt-4.1-nano\")\n", + "\n", + "# Create the writer agent\n", + "writer = AssistantAgent(\"writer\", model_client=client, system_message=\"Draft a short paragraph on climate change.\")\n", + "\n", + "# Create the reviewer agent\n", + "reviewer = AssistantAgent(\"reviewer\", model_client=client, system_message=\"Review the draft and suggest improvements.\")\n", + "\n", + "# Build the graph\n", + "builder = DiGraphBuilder()\n", + "builder.add_node(writer).add_node(reviewer)\n", + "builder.add_edge(writer, reviewer)\n", + "\n", + "# Build and validate the graph\n", + "graph = builder.build()\n", + "\n", + "# Create the flow\n", + "flow = GraphFlow([writer, reviewer], graph=graph)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1b19f9a6", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "source='user' models_usage=None metadata={} content='Write a short paragraph about climate change.' type='TextMessage'\n", + "source='writer' models_usage=RequestUsage(prompt_tokens=28, completion_tokens=95) metadata={} content='Climate change refers to long-term shifts in temperature, precipitation, and other atmospheric patterns, largely driven by human activities such as burning fossil fuels, deforestation, and industrial processes. These changes contribute to rising global temperatures, melting ice caps, more frequent and severe weather events, and adverse impacts on ecosystems and human communities. Addressing climate change requires global cooperation to reduce greenhouse gas emissions, transition to renewable energy sources, and implement sustainable practices to protect the planet for future generations.' type='TextMessage'\n", + "source='reviewer' models_usage=RequestUsage(prompt_tokens=127, completion_tokens=144) metadata={} content=\"The paragraph provides a clear overview of climate change, its causes, and its impacts. To enhance clarity and engagement, consider adding specific examples or emphasizing the urgency of action. Here's a revised version:\\n\\nClimate change is a long-term alteration of Earth's climate patterns caused primarily by human activities such as burning fossil fuels, deforestation, and industrial emissions. These actions increase greenhouse gases in the atmosphere, leading to rising global temperatures, melting ice caps, and more frequent extreme weather events like hurricanes and droughts. The effects threaten ecosystems, disrupt agriculture, and endanger communities worldwide. Addressing this crisis requires urgent, coordinated global efforts to reduce emissions, adopt renewable energy, and promote sustainable practices to safeguard the planet for future generations.\" type='TextMessage'\n", + "source='DiGraphStopAgent' models_usage=None metadata={} content='Digraph execution is complete' type='StopMessage'\n", + "messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Write a short paragraph about climate change.', type='TextMessage'), TextMessage(source='writer', models_usage=RequestUsage(prompt_tokens=28, completion_tokens=95), metadata={}, content='Climate change refers to long-term shifts in temperature, precipitation, and other atmospheric patterns, largely driven by human activities such as burning fossil fuels, deforestation, and industrial processes. These changes contribute to rising global temperatures, melting ice caps, more frequent and severe weather events, and adverse impacts on ecosystems and human communities. Addressing climate change requires global cooperation to reduce greenhouse gas emissions, transition to renewable energy sources, and implement sustainable practices to protect the planet for future generations.', type='TextMessage'), TextMessage(source='reviewer', models_usage=RequestUsage(prompt_tokens=127, completion_tokens=144), metadata={}, content=\"The paragraph provides a clear overview of climate change, its causes, and its impacts. To enhance clarity and engagement, consider adding specific examples or emphasizing the urgency of action. Here's a revised version:\\n\\nClimate change is a long-term alteration of Earth's climate patterns caused primarily by human activities such as burning fossil fuels, deforestation, and industrial emissions. These actions increase greenhouse gases in the atmosphere, leading to rising global temperatures, melting ice caps, and more frequent extreme weather events like hurricanes and droughts. The effects threaten ecosystems, disrupt agriculture, and endanger communities worldwide. Addressing this crisis requires urgent, coordinated global efforts to reduce emissions, adopt renewable energy, and promote sustainable practices to safeguard the planet for future generations.\", type='TextMessage'), StopMessage(source='DiGraphStopAgent', models_usage=None, metadata={}, content='Digraph execution is complete', type='StopMessage')] stop_reason='Stop message received'\n" + ] + } + ], + "source": [ + "# Use `asyncio.run(...)` and wrap the below in a async function when running in a script.\n", + "stream = flow.run_stream(task=\"Write a short paragraph about climate change.\")\n", + "async for event in stream: # type: ignore\n", + " print(event)\n", + "# Use Console(flow.run_stream(...)) for better formatting in console." + ] + }, + { + "cell_type": "markdown", + "id": "5157cfbf", + "metadata": {}, + "source": [ + "### Parallel Flow with Join\n", + "\n", + "We now create a slightly more complex flow:\n", + "\n", + "- A **writer** drafts a paragraph.\n", + "- Two **editors** independently edit for grammar and style (parallel fan-out).\n", + "- A **final reviewer** consolidates their edits (join).\n", + "\n", + "Execution starts at the **writer**, fans out to **editor1** and **editor2** simultaneously, and then both feed into the **final reviewer**.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a54d2454", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- TextMessage (user) ----------\n", + "Write a short paragraph about climate change.\n", + "---------- TextMessage (writer) ----------\n", + "Climate change refers to long-term shifts in weather patterns and global temperatures, largely driven by human activities such as burning fossil fuels, deforestation, and industrial processes. These activities increase concentrations of greenhouse gases like carbon dioxide and methane in the atmosphere, leading to global warming. The impacts of climate change include more frequent and severe weather events, rising sea levels, melting glaciers, and disruptions to ecosystems and agriculture. Addressing this urgent issue requires international cooperation, significant shifts toward renewable energy sources, and sustainable practices to reduce our carbon footprint and protect the planet for future generations.\n", + "---------- TextMessage (editor1) ----------\n", + "Climate change refers to long-term shifts in weather patterns and global temperatures, largely driven by human activities such as burning fossil fuels, deforestation, and industrial processes. These activities increase concentrations of greenhouse gases like carbon dioxide and methane in the atmosphere, leading to global warming. The impacts of climate change include more frequent and severe weather events, rising sea levels, melting glaciers, and disruptions to ecosystems and agriculture. Addressing this urgent issue requires international cooperation, significant shifts toward renewable energy sources, and sustainable practices to reduce our carbon footprint and protect the planet for future generations.\n", + "---------- TextMessage (editor2) ----------\n", + "Climate change involves long-term alterations in weather patterns and global temperatures, primarily caused by human activities like burning fossil fuels, deforestation, and industrial processes. These actions elevate levels of greenhouse gases such as carbon dioxide and methane, resulting in global warming. Its consequences are widespread, including more frequent and intense storms, rising sea levels, melting glaciers, and disturbances to ecosystems and agriculture. Combating this crisis demands international collaboration, a swift transition to renewable energy, and sustainable practices to cut carbon emissions, ensuring a healthier planet for future generations.\n", + "---------- TextMessage (final_reviewer) ----------\n", + "Climate change involves long-term alterations in weather patterns and global temperatures, primarily caused by human activities such as burning fossil fuels, deforestation, and industrial processes. These actions increase levels of greenhouse gases like carbon dioxide and methane, leading to global warming. Its consequences include more frequent and intense storms, rising sea levels, melting glaciers, and disruptions to ecosystems and agriculture. Addressing this crisis requires international collaboration, a swift transition to renewable energy, and sustainable practices to reduce carbon emissions, ensuring a healthier planet for future generations.\n", + "---------- StopMessage (DiGraphStopAgent) ----------\n", + "Digraph execution is complete\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Write a short paragraph about climate change.', type='TextMessage'), TextMessage(source='writer', models_usage=RequestUsage(prompt_tokens=28, completion_tokens=113), metadata={}, content='Climate change refers to long-term shifts in weather patterns and global temperatures, largely driven by human activities such as burning fossil fuels, deforestation, and industrial processes. These activities increase concentrations of greenhouse gases like carbon dioxide and methane in the atmosphere, leading to global warming. The impacts of climate change include more frequent and severe weather events, rising sea levels, melting glaciers, and disruptions to ecosystems and agriculture. Addressing this urgent issue requires international cooperation, significant shifts toward renewable energy sources, and sustainable practices to reduce our carbon footprint and protect the planet for future generations.', type='TextMessage'), TextMessage(source='editor1', models_usage=RequestUsage(prompt_tokens=144, completion_tokens=113), metadata={}, content='Climate change refers to long-term shifts in weather patterns and global temperatures, largely driven by human activities such as burning fossil fuels, deforestation, and industrial processes. These activities increase concentrations of greenhouse gases like carbon dioxide and methane in the atmosphere, leading to global warming. The impacts of climate change include more frequent and severe weather events, rising sea levels, melting glaciers, and disruptions to ecosystems and agriculture. Addressing this urgent issue requires international cooperation, significant shifts toward renewable energy sources, and sustainable practices to reduce our carbon footprint and protect the planet for future generations.', type='TextMessage'), TextMessage(source='editor2', models_usage=RequestUsage(prompt_tokens=263, completion_tokens=107), metadata={}, content='Climate change involves long-term alterations in weather patterns and global temperatures, primarily caused by human activities like burning fossil fuels, deforestation, and industrial processes. These actions elevate levels of greenhouse gases such as carbon dioxide and methane, resulting in global warming. Its consequences are widespread, including more frequent and intense storms, rising sea levels, melting glaciers, and disturbances to ecosystems and agriculture. Combating this crisis demands international collaboration, a swift transition to renewable energy, and sustainable practices to cut carbon emissions, ensuring a healthier planet for future generations.', type='TextMessage'), TextMessage(source='final_reviewer', models_usage=RequestUsage(prompt_tokens=383, completion_tokens=104), metadata={}, content='Climate change involves long-term alterations in weather patterns and global temperatures, primarily caused by human activities such as burning fossil fuels, deforestation, and industrial processes. These actions increase levels of greenhouse gases like carbon dioxide and methane, leading to global warming. Its consequences include more frequent and intense storms, rising sea levels, melting glaciers, and disruptions to ecosystems and agriculture. Addressing this crisis requires international collaboration, a swift transition to renewable energy, and sustainable practices to reduce carbon emissions, ensuring a healthier planet for future generations.', type='TextMessage'), StopMessage(source='DiGraphStopAgent', models_usage=None, metadata={}, content='Digraph execution is complete', type='StopMessage')], stop_reason='Stop message received')" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.teams import DiGraphBuilder, GraphFlow\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "# Create an OpenAI model client\n", + "client = OpenAIChatCompletionClient(model=\"gpt-4.1-nano\")\n", + "\n", + "# Create the writer agent\n", + "writer = AssistantAgent(\"writer\", model_client=client, system_message=\"Draft a short paragraph on climate change.\")\n", + "\n", + "# Create two editor agents\n", + "editor1 = AssistantAgent(\"editor1\", model_client=client, system_message=\"Edit the paragraph for grammar.\")\n", + "\n", + "editor2 = AssistantAgent(\"editor2\", model_client=client, system_message=\"Edit the paragraph for style.\")\n", + "\n", + "# Create the final reviewer agent\n", + "final_reviewer = AssistantAgent(\n", + " \"final_reviewer\",\n", + " model_client=client,\n", + " system_message=\"Consolidate the grammar and style edits into a final version.\",\n", + ")\n", + "\n", + "# Build the workflow graph\n", + "builder = DiGraphBuilder()\n", + "builder.add_node(writer).add_node(editor1).add_node(editor2).add_node(final_reviewer)\n", + "\n", + "# Fan-out from writer to editor1 and editor2\n", + "builder.add_edge(writer, editor1)\n", + "builder.add_edge(writer, editor2)\n", + "\n", + "# Fan-in both editors into final reviewer\n", + "builder.add_edge(editor1, final_reviewer)\n", + "builder.add_edge(editor2, final_reviewer)\n", + "\n", + "# Build and validate the graph\n", + "graph = builder.build()\n", + "\n", + "# Create the flow\n", + "flow = GraphFlow(\n", + " participants=builder.get_participants(),\n", + " graph=graph,\n", + ")\n", + "\n", + "# Run the workflow\n", + "await Console(flow.run_stream(task=\"Write a short paragraph about climate change.\"))" + ] + }, + { + "cell_type": "markdown", + "id": "0343182c", + "metadata": {}, + "source": [ + "## Message Filtering\n", + "\n", + "### Execution Graph vs. Message Graph\n", + "\n", + "In {py:class}`~autogen_agentchat.teams.GraphFlow`, the **execution graph** is defined using \n", + "{py:class}`~autogen_agentchat.teams.DiGraph`, which controls the order in which agents execute.\n", + "However, the execution graph does not control what messages an agent receives from other agents.\n", + "By default, all messages are sent to all agents in the graph.\n", + "\n", + "**Message filtering** is a separate feature that allows you to filter the messages\n", + "received by each agent and limiting their model context to only the relevant information.\n", + "The set of message filters defines the **message graph** in the flow.\n", + "\n", + "Specifying the message graph can help with:\n", + "- Reduce hallucinations\n", + "- Control memory load\n", + "- Focus agents only on relevant information\n", + "\n", + "You can use {py:class}`~autogen_agentchat.agents.MessageFilterAgent` together with {py:class}`~autogen_agentchat.agents.MessageFilterConfig` and {py:class}`~autogen_agentchat.agents.PerSourceFilter` to define these rules." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2a59af03", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- TextMessage (user) ----------\n", + "Summarize key facts about climate change.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- TextMessage (researcher) ----------\n", + "Certainly! Here are some key facts about climate change:\n", + "\n", + "1. **Global Warming**: Earth's average surface temperature has increased significantly over the past century, primarily due to human activities.\n", + "2. **Greenhouse Gas Emissions**: The main contributors are carbon dioxide (CO₂), methane (CH₄), and nitrous oxide (N₂O), resulting from burning fossil fuels, deforestation, and industrial processes.\n", + "3. **Impacts on Weather and Climate**: Climate change leads to more frequent and severe heatwaves, storms, droughts, and heavy rainfall.\n", + "4. **Rising Sea Levels**: Melting polar ice caps and glaciers, along with thermal expansion of seawater, are causing sea levels to rise.\n", + "5. **Effects on Ecosystems**: Altered habitats threaten plant and animal species, leading to biodiversity loss.\n", + "6. **Human Health and Societies**: Climate change contributes to health issues, food and water insecurity, and displacement of populations.\n", + "7. **Global Response**: International efforts like the Paris Agreement aim to limit temperature rise, promote renewable energy, and reduce emissions.\n", + "8. **Urgency**: Addressing climate change requires immediate, concerted actions to mitigate further damage and adapt to changes.\n", + "\n", + "Let me know if you want more detailed information on any of these points!\n", + "---------- TextMessage (analyst) ----------\n", + "Your summary effectively covers the fundamental aspects of climate change and presents them clearly. Here are some suggestions to improve clarity, depth, and engagement:\n", + "\n", + "1. Enhance structure with subheadings: Organize points into thematic sections (e.g., Causes, Effects, Responses) for easier navigation.\n", + "2. Add recent context or data: Incorporate the latest statistics or notable recent events to emphasize urgency.\n", + "3. Emphasize solutions: Briefly mention specific mitigation and adaptation strategies beyond international agreements.\n", + "4. Use more precise language: For example, specify the amount of temperature increase globally (~1.2°C since pre-industrial times).\n", + "5. Incorporate the importance of individual actions: Highlight how personal choices contribute to climate efforts.\n", + "6. Mention climate feedback loops: Briefly note how certain effects (like melting ice) can accelerate warming.\n", + "\n", + "**Improved Version:**\n", + "\n", + "---\n", + "\n", + "**Overview of Climate Change**\n", + "\n", + "**Causes:**\n", + "- Human activities, especially burning fossil fuels, deforestation, and industrial processes, have led to increased concentrations of greenhouse gases such as carbon dioxide (CO₂), methane (CH₄), and nitrous oxide (N₂O).\n", + "- Since the late 19th century, Earth's average surface temperature has risen by approximately 1.2°C, with the past decade being the warmest on record.\n", + "\n", + "**Impacts:**\n", + "- The changing climate causes more frequent and intense heatwaves, storms, droughts, and heavy rainfall events.\n", + "- Melting polar ice caps and glaciers, along with thermal expansion, are raising sea levels, threatening coastal communities.\n", + "- Ecosystems are shifting, leading to habitat loss and risking biodiversity, with some species facing extinction.\n", + "- Human health and societies are affected through increased heat-related illnesses, food and water insecurity, and displacement due to extreme weather events.\n", + "\n", + "**Global Response and Solutions:**\n", + "- International agreements like the Paris Agreement aim to limit global temperature rise well below 2°C.\n", + "- Strategies include transitioning to renewable energy sources, increasing energy efficiency, reforestation, and sustainable land use.\n", + "- Community and individual actions—reducing carbon footprints, supporting sustainable policies, and raising awareness—are essential components.\n", + "\n", + "**Urgency and Call to Action:**\n", + "- Immediate, coordinated efforts are critical to mitigate irreversible damage and adapt to ongoing changes.\n", + "- Every sector, from government to individual, has a role to play in creating a sustainable future.\n", + "\n", + "---\n", + "\n", + "Let me know if you'd like a more detailed explanation of any section or additional statistical data!\n", + "---------- TextMessage (presenter) ----------\n", + "**Slide Title:** \n", + "**Climate Change: Causes, Impacts & Solutions**\n", + "\n", + "**Causes:** \n", + "- Emissions from burning fossil fuels, deforestation, industrial activities \n", + "- Greenhouse gases (CO₂, CH₄, N₂O) have increased significantly \n", + "- Global temperature has risen by ~1.2°C since pre-industrial times \n", + "\n", + "**Impacts:** \n", + "- More frequent heatwaves, storms, droughts, and heavy rainfall \n", + "- Melting ice caps and rising sea levels threaten coastal areas \n", + "- Habitat loss and decreased biodiversity \n", + "- Health risks and societal disruptions \n", + "\n", + "**Responses & Solutions:** \n", + "- International efforts like the Paris Agreement aim to limit warming \n", + "- Transitioning to renewable energy, energy efficiency, reforestation \n", + "- Community and individual actions: reducing carbon footprints and raising awareness \n", + "\n", + "**Urgency:** \n", + "- Immediate, coordinated action is essential to prevent irreversible damage \n", + "- Everyone has a role in building a sustainable future\n", + "---------- StopMessage (DiGraphStopAgent) ----------\n", + "Digraph execution is complete\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Summarize key facts about climate change.', type='TextMessage'), TextMessage(source='researcher', models_usage=RequestUsage(prompt_tokens=30, completion_tokens=267), metadata={}, content=\"Certainly! Here are some key facts about climate change:\\n\\n1. **Global Warming**: Earth's average surface temperature has increased significantly over the past century, primarily due to human activities.\\n2. **Greenhouse Gas Emissions**: The main contributors are carbon dioxide (CO₂), methane (CH₄), and nitrous oxide (N₂O), resulting from burning fossil fuels, deforestation, and industrial processes.\\n3. **Impacts on Weather and Climate**: Climate change leads to more frequent and severe heatwaves, storms, droughts, and heavy rainfall.\\n4. **Rising Sea Levels**: Melting polar ice caps and glaciers, along with thermal expansion of seawater, are causing sea levels to rise.\\n5. **Effects on Ecosystems**: Altered habitats threaten plant and animal species, leading to biodiversity loss.\\n6. **Human Health and Societies**: Climate change contributes to health issues, food and water insecurity, and displacement of populations.\\n7. **Global Response**: International efforts like the Paris Agreement aim to limit temperature rise, promote renewable energy, and reduce emissions.\\n8. **Urgency**: Addressing climate change requires immediate, concerted actions to mitigate further damage and adapt to changes.\\n\\nLet me know if you want more detailed information on any of these points!\", type='TextMessage'), TextMessage(source='analyst', models_usage=RequestUsage(prompt_tokens=287, completion_tokens=498), metadata={}, content=\"Your summary effectively covers the fundamental aspects of climate change and presents them clearly. Here are some suggestions to improve clarity, depth, and engagement:\\n\\n1. Enhance structure with subheadings: Organize points into thematic sections (e.g., Causes, Effects, Responses) for easier navigation.\\n2. Add recent context or data: Incorporate the latest statistics or notable recent events to emphasize urgency.\\n3. Emphasize solutions: Briefly mention specific mitigation and adaptation strategies beyond international agreements.\\n4. Use more precise language: For example, specify the amount of temperature increase globally (~1.2°C since pre-industrial times).\\n5. Incorporate the importance of individual actions: Highlight how personal choices contribute to climate efforts.\\n6. Mention climate feedback loops: Briefly note how certain effects (like melting ice) can accelerate warming.\\n\\n**Improved Version:**\\n\\n---\\n\\n**Overview of Climate Change**\\n\\n**Causes:**\\n- Human activities, especially burning fossil fuels, deforestation, and industrial processes, have led to increased concentrations of greenhouse gases such as carbon dioxide (CO₂), methane (CH₄), and nitrous oxide (N₂O).\\n- Since the late 19th century, Earth's average surface temperature has risen by approximately 1.2°C, with the past decade being the warmest on record.\\n\\n**Impacts:**\\n- The changing climate causes more frequent and intense heatwaves, storms, droughts, and heavy rainfall events.\\n- Melting polar ice caps and glaciers, along with thermal expansion, are raising sea levels, threatening coastal communities.\\n- Ecosystems are shifting, leading to habitat loss and risking biodiversity, with some species facing extinction.\\n- Human health and societies are affected through increased heat-related illnesses, food and water insecurity, and displacement due to extreme weather events.\\n\\n**Global Response and Solutions:**\\n- International agreements like the Paris Agreement aim to limit global temperature rise well below 2°C.\\n- Strategies include transitioning to renewable energy sources, increasing energy efficiency, reforestation, and sustainable land use.\\n- Community and individual actions—reducing carbon footprints, supporting sustainable policies, and raising awareness—are essential components.\\n\\n**Urgency and Call to Action:**\\n- Immediate, coordinated efforts are critical to mitigate irreversible damage and adapt to ongoing changes.\\n- Every sector, from government to individual, has a role to play in creating a sustainable future.\\n\\n---\\n\\nLet me know if you'd like a more detailed explanation of any section or additional statistical data!\", type='TextMessage'), TextMessage(source='presenter', models_usage=RequestUsage(prompt_tokens=521, completion_tokens=192), metadata={}, content='**Slide Title:** \\n**Climate Change: Causes, Impacts & Solutions**\\n\\n**Causes:** \\n- Emissions from burning fossil fuels, deforestation, industrial activities \\n- Greenhouse gases (CO₂, CH₄, N₂O) have increased significantly \\n- Global temperature has risen by ~1.2°C since pre-industrial times \\n\\n**Impacts:** \\n- More frequent heatwaves, storms, droughts, and heavy rainfall \\n- Melting ice caps and rising sea levels threaten coastal areas \\n- Habitat loss and decreased biodiversity \\n- Health risks and societal disruptions \\n\\n**Responses & Solutions:** \\n- International efforts like the Paris Agreement aim to limit warming \\n- Transitioning to renewable energy, energy efficiency, reforestation \\n- Community and individual actions: reducing carbon footprints and raising awareness \\n\\n**Urgency:** \\n- Immediate, coordinated action is essential to prevent irreversible damage \\n- Everyone has a role in building a sustainable future', type='TextMessage'), StopMessage(source='DiGraphStopAgent', models_usage=None, metadata={}, content='Digraph execution is complete', type='StopMessage')], stop_reason='Stop message received')" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from autogen_agentchat.agents import AssistantAgent, MessageFilterAgent, MessageFilterConfig, PerSourceFilter\n", + "from autogen_agentchat.teams import DiGraphBuilder, GraphFlow\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "# Model client\n", + "client = OpenAIChatCompletionClient(model=\"gpt-4.1-nano\")\n", + "\n", + "# Create agents\n", + "researcher = AssistantAgent(\n", + " \"researcher\", model_client=client, system_message=\"Summarize key facts about climate change.\"\n", + ")\n", + "analyst = AssistantAgent(\"analyst\", model_client=client, system_message=\"Review the summary and suggest improvements.\")\n", + "presenter = AssistantAgent(\n", + " \"presenter\", model_client=client, system_message=\"Prepare a presentation slide based on the final summary.\"\n", + ")\n", + "\n", + "# Apply message filtering\n", + "filtered_analyst = MessageFilterAgent(\n", + " name=\"analyst\",\n", + " wrapped_agent=analyst,\n", + " filter=MessageFilterConfig(per_source=[PerSourceFilter(source=\"researcher\", position=\"last\", count=1)]),\n", + ")\n", + "\n", + "filtered_presenter = MessageFilterAgent(\n", + " name=\"presenter\",\n", + " wrapped_agent=presenter,\n", + " filter=MessageFilterConfig(per_source=[PerSourceFilter(source=\"analyst\", position=\"last\", count=1)]),\n", + ")\n", + "\n", + "# Build the flow\n", + "builder = DiGraphBuilder()\n", + "builder.add_node(researcher).add_node(filtered_analyst).add_node(filtered_presenter)\n", + "builder.add_edge(researcher, filtered_analyst).add_edge(filtered_analyst, filtered_presenter)\n", + "\n", + "# Create the flow\n", + "flow = GraphFlow(\n", + " participants=builder.get_participants(),\n", + " graph=builder.build(),\n", + ")\n", + "\n", + "# Run the flow\n", + "await Console(flow.run_stream(task=\"Summarize key facts about climate change.\"))" + ] + }, + { + "cell_type": "markdown", + "id": "f7309529", + "metadata": {}, + "source": [ + "## 🔁 Advanced Example: Conditional Loop + Filtered Summary\n", + "\n", + "This example demonstrates:\n", + "\n", + "- A loop between generator and reviewer (which exits when reviewer says \"APPROVE\")\n", + "- A summarizer agent that only sees the first user input and the last reviewer message\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "af297db2", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- TextMessage (user) ----------\n", + "Brainstorm ways to reduce plastic waste.\n", + "---------- TextMessage (generator) ----------\n", + "Here are some creative ideas to help reduce plastic waste:\n", + "\n", + "1. **Refill Stations**: Create refill stations for common household liquids (like soaps, shampoos, and detergents) where people can bring their own containers to fill up.\n", + "\n", + "2. **DIY Kits**: Offer DIY kits for making eco-friendly products, such as beeswax wraps, reusable bags, or natural cleaning solutions.\n", + "\n", + "3. **Community Swap Events**: Organize swap events where people can bring unwanted items, including clothing and household products, to exchange instead of purchasing new items.\n", + "\n", + "4. **Plastic-Free Challenge**: Launch a community-wide challenge encouraging residents to go plastic-free for a month, sharing tips and experiences on social media.\n", + "\n", + "5. **Incentivize Businesses**: Create incentives for local businesses that implement sustainable practices, like providing discounts to customers who bring their own containers or bags.\n", + "\n", + "6. **Educational Campaigns**: Partner with schools to educate children about the impact of plastic waste and encourage them to take home messages to their families.\n", + "\n", + "7. **Plastic-Free Shopping Zones**: Designate certain areas in town as plastic-free zones where businesses agree to eliminate single-use plastics.\n", + "\n", + "8. **Upcycling Workshops**: Host workshops teaching people how to upcycle plastic waste into art, furniture, or home decor.\n", + "\n", + "9. **Composting Competition**: Encourage households to compost food waste and offer a competition for the best composting garden to foster eco-awareness.\n", + "\n", + "10. **Zero-Waste Stores**: Support or start zero-waste stores that sell bulk goods, allowing customers to shop with their own containers.\n", + "\n", + "11. **Mobile Recycling Units**: Implement mobile recycling units that visit neighborhoods to educate residents on recycling properly and collect recyclables.\n", + "\n", + "12. **Plastic Offset Programs**: Create programs that allow individuals and companies to offset their plastic usage through donations to initiatives that remove plastic from oceans.\n", + "\n", + "13. **Collaboration with Influencers**: Partner with social media influencers to spread the message about reducing plastic waste and promote sustainable alternatives.\n", + "\n", + "14. **Sustainable Product Market**: Organize a market dedicated exclusively to sustainable products and services, showcasing local vendors who prioritize eco-friendly practices.\n", + "\n", + "15. **Plastic Waste Art Installations**: Collaborate with artists to create public installations made from recycled plastic, raising awareness of the plastic problem in a visually impactful way.\n", + "\n", + "16. **Interactive Apps**: Develop apps that track plastic usage and provide personalized tips for reducing plastic consumption based on user habits.\n", + "\n", + "17. **Corporate Partnerships**: Work with businesses to develop corporate responsibility programs focused on reducing plastic use in their operations and packaging.\n", + "\n", + "18. **Legislation Advocacy**: Promote local policies that restrict single-use plastics or support more effective recycling programs.\n", + "\n", + "19. **Public Transportation Awareness**: Encourage public transportation usage by providing eco-friendly incentives for those who walk, bike, or use buses instead of cars.\n", + "\n", + "20. **Create a Local Plastic Waste Repository**: Start a community hub where individuals and artists can drop off plastic waste for reuse in projects or art pieces.\n", + "\n", + "By implementing these ideas, communities can take significant steps toward reducing plastic waste and fostering a sustainable future.\n", + "---------- TextMessage (reviewer) ----------\n", + "These ideas present a comprehensive and practical approach to reducing plastic waste within communities. Here’s some feedback and considerations for each suggestion:\n", + "\n", + "1. **Refill Stations**: Great idea; consider partnering with local health and wellness shops for broader adoption.\n", + " \n", + "2. **DIY Kits**: Ensure kits include clear instructions and safety guidance to promote user-friendliness.\n", + "\n", + "3. **Community Swap Events**: Promote these as regular events to build a sense of community and reinforce sustainable habits.\n", + "\n", + "4. **Plastic-Free Challenge**: Consider creating a dedicated hashtag to track participants’ journeys and foster engagement online.\n", + "\n", + "5. **Incentivize Businesses**: Work on a simple certification system for sustainable businesses to encourage participation and recognition.\n", + "\n", + "6. **Educational Campaigns**: Tailor content to different age groups to maximize impact across the community.\n", + "\n", + "7. **Plastic-Free Shopping Zones**: Consider involving local government for support and promotion, which can increase visibility and compliance.\n", + "\n", + "8. **Upcycling Workshops**: Source materials locally for workshops to decrease transportation emissions and support local businesses.\n", + "\n", + "9. **Composting Competition**: Collaborate with gardening clubs for expert insights and to broaden participation in community gardening.\n", + "\n", + "10. **Zero-Waste Stores**: Explore online sales options to enhance accessibility while retaining the focus on zero-waste practices.\n", + "\n", + "11. **Mobile Recycling Units**: Train volunteers for effective community engagement and education during visits.\n", + "\n", + "12. **Plastic Offset Programs**: Emphasize transparency in how donations are used to build trust within the community.\n", + "\n", + "13. **Collaboration with Influencers**: Ensure influencers have a genuine commitment to sustainability to ensure credible messaging.\n", + "\n", + "14. **Sustainable Product Market**: Regularly invite new vendors to keep the market fresh and encourage innovation in sustainable products.\n", + "\n", + "15. **Plastic Waste Art Installations**: Consider educational plaques accompanying installations that inform viewers about the issues of plastic waste.\n", + "\n", + "16. **Interactive Apps**: Include gamification elements to encourage increased user engagement and sharing among friends.\n", + "\n", + "17. **Corporate Partnerships**: Develop case studies or success stories to showcase the benefits of reduced plastic use for businesses.\n", + "\n", + "18. **Legislation Advocacy**: Mobilize community members to become advocates themselves, creating a grass-roots effort for policy change.\n", + "\n", + "19. **Public Transportation Awareness**: Explore partnerships with public transit systems for eco-friendly promotions.\n", + "\n", + "20. **Create a Local Plastic Waste Repository**: Establish partnerships with local craft schools or organizations to enhance creativity and use of the repository.\n", + "\n", + "Overall, these ideas have high potential for impactful implementation. Emphasizing community engagement, education, and ongoing support will help ensure their success. **APPROVE**.\n", + "---------- TextMessage (summary) ----------\n", + "The user requested brainstorming ideas to reduce plastic waste, looking for practical and impactful solutions. The reviewer provided detailed feedback on each suggested idea, indicating strengths and considerations for improvement, such as involving local businesses, enhancing community engagement, and promoting educational initiatives. The final feedback indicates that the suggestions have great potential and emphasizes the importance of community involvement and education for successful implementation, culminating in an overall approval of the ideas presented.\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(id='eca90b4f-a8cc-4f06-9b42-d8387caf338e', source='user', models_usage=None, metadata={}, created_at=datetime.datetime(2025, 7, 15, 1, 48, 51, 648989, tzinfo=datetime.timezone.utc), content='Brainstorm ways to reduce plastic waste.', type='TextMessage'), TextMessage(id='29767cbd-ae8d-4dfb-be57-7f982aaddc4b', source='generator', models_usage=RequestUsage(prompt_tokens=27, completion_tokens=627), metadata={}, created_at=datetime.datetime(2025, 7, 15, 1, 49, 6, 788238, tzinfo=datetime.timezone.utc), content='Here are some creative ideas to help reduce plastic waste:\\n\\n1. **Refill Stations**: Create refill stations for common household liquids (like soaps, shampoos, and detergents) where people can bring their own containers to fill up.\\n\\n2. **DIY Kits**: Offer DIY kits for making eco-friendly products, such as beeswax wraps, reusable bags, or natural cleaning solutions.\\n\\n3. **Community Swap Events**: Organize swap events where people can bring unwanted items, including clothing and household products, to exchange instead of purchasing new items.\\n\\n4. **Plastic-Free Challenge**: Launch a community-wide challenge encouraging residents to go plastic-free for a month, sharing tips and experiences on social media.\\n\\n5. **Incentivize Businesses**: Create incentives for local businesses that implement sustainable practices, like providing discounts to customers who bring their own containers or bags.\\n\\n6. **Educational Campaigns**: Partner with schools to educate children about the impact of plastic waste and encourage them to take home messages to their families.\\n\\n7. **Plastic-Free Shopping Zones**: Designate certain areas in town as plastic-free zones where businesses agree to eliminate single-use plastics.\\n\\n8. **Upcycling Workshops**: Host workshops teaching people how to upcycle plastic waste into art, furniture, or home decor.\\n\\n9. **Composting Competition**: Encourage households to compost food waste and offer a competition for the best composting garden to foster eco-awareness.\\n\\n10. **Zero-Waste Stores**: Support or start zero-waste stores that sell bulk goods, allowing customers to shop with their own containers.\\n\\n11. **Mobile Recycling Units**: Implement mobile recycling units that visit neighborhoods to educate residents on recycling properly and collect recyclables.\\n\\n12. **Plastic Offset Programs**: Create programs that allow individuals and companies to offset their plastic usage through donations to initiatives that remove plastic from oceans.\\n\\n13. **Collaboration with Influencers**: Partner with social media influencers to spread the message about reducing plastic waste and promote sustainable alternatives.\\n\\n14. **Sustainable Product Market**: Organize a market dedicated exclusively to sustainable products and services, showcasing local vendors who prioritize eco-friendly practices.\\n\\n15. **Plastic Waste Art Installations**: Collaborate with artists to create public installations made from recycled plastic, raising awareness of the plastic problem in a visually impactful way.\\n\\n16. **Interactive Apps**: Develop apps that track plastic usage and provide personalized tips for reducing plastic consumption based on user habits.\\n\\n17. **Corporate Partnerships**: Work with businesses to develop corporate responsibility programs focused on reducing plastic use in their operations and packaging.\\n\\n18. **Legislation Advocacy**: Promote local policies that restrict single-use plastics or support more effective recycling programs.\\n\\n19. **Public Transportation Awareness**: Encourage public transportation usage by providing eco-friendly incentives for those who walk, bike, or use buses instead of cars.\\n\\n20. **Create a Local Plastic Waste Repository**: Start a community hub where individuals and artists can drop off plastic waste for reuse in projects or art pieces.\\n\\nBy implementing these ideas, communities can take significant steps toward reducing plastic waste and fostering a sustainable future.', type='TextMessage'), TextMessage(id='54e02028-0239-4809-8163-af60745e6b9d', source='reviewer', models_usage=RequestUsage(prompt_tokens=671, completion_tokens=532), metadata={}, created_at=datetime.datetime(2025, 7, 15, 1, 49, 17, 327641, tzinfo=datetime.timezone.utc), content='These ideas present a comprehensive and practical approach to reducing plastic waste within communities. Here’s some feedback and considerations for each suggestion:\\n\\n1. **Refill Stations**: Great idea; consider partnering with local health and wellness shops for broader adoption.\\n \\n2. **DIY Kits**: Ensure kits include clear instructions and safety guidance to promote user-friendliness.\\n\\n3. **Community Swap Events**: Promote these as regular events to build a sense of community and reinforce sustainable habits.\\n\\n4. **Plastic-Free Challenge**: Consider creating a dedicated hashtag to track participants’ journeys and foster engagement online.\\n\\n5. **Incentivize Businesses**: Work on a simple certification system for sustainable businesses to encourage participation and recognition.\\n\\n6. **Educational Campaigns**: Tailor content to different age groups to maximize impact across the community.\\n\\n7. **Plastic-Free Shopping Zones**: Consider involving local government for support and promotion, which can increase visibility and compliance.\\n\\n8. **Upcycling Workshops**: Source materials locally for workshops to decrease transportation emissions and support local businesses.\\n\\n9. **Composting Competition**: Collaborate with gardening clubs for expert insights and to broaden participation in community gardening.\\n\\n10. **Zero-Waste Stores**: Explore online sales options to enhance accessibility while retaining the focus on zero-waste practices.\\n\\n11. **Mobile Recycling Units**: Train volunteers for effective community engagement and education during visits.\\n\\n12. **Plastic Offset Programs**: Emphasize transparency in how donations are used to build trust within the community.\\n\\n13. **Collaboration with Influencers**: Ensure influencers have a genuine commitment to sustainability to ensure credible messaging.\\n\\n14. **Sustainable Product Market**: Regularly invite new vendors to keep the market fresh and encourage innovation in sustainable products.\\n\\n15. **Plastic Waste Art Installations**: Consider educational plaques accompanying installations that inform viewers about the issues of plastic waste.\\n\\n16. **Interactive Apps**: Include gamification elements to encourage increased user engagement and sharing among friends.\\n\\n17. **Corporate Partnerships**: Develop case studies or success stories to showcase the benefits of reduced plastic use for businesses.\\n\\n18. **Legislation Advocacy**: Mobilize community members to become advocates themselves, creating a grass-roots effort for policy change.\\n\\n19. **Public Transportation Awareness**: Explore partnerships with public transit systems for eco-friendly promotions.\\n\\n20. **Create a Local Plastic Waste Repository**: Establish partnerships with local craft schools or organizations to enhance creativity and use of the repository.\\n\\nOverall, these ideas have high potential for impactful implementation. Emphasizing community engagement, education, and ongoing support will help ensure their success. **APPROVE**.', type='TextMessage'), TextMessage(id='55409dc3-9766-4071-ab85-0b3125cb59c7', source='summary', models_usage=RequestUsage(prompt_tokens=570, completion_tokens=82), metadata={}, created_at=datetime.datetime(2025, 7, 15, 1, 49, 19, 442276, tzinfo=datetime.timezone.utc), content='The user requested brainstorming ideas to reduce plastic waste, looking for practical and impactful solutions. The reviewer provided detailed feedback on each suggested idea, indicating strengths and considerations for improvement, such as involving local businesses, enhancing community engagement, and promoting educational initiatives. The final feedback indicates that the suggestions have great potential and emphasizes the importance of community involvement and education for successful implementation, culminating in an overall approval of the ideas presented.', type='TextMessage')], stop_reason='Digraph execution is complete')" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from autogen_agentchat.agents import AssistantAgent, MessageFilterAgent, MessageFilterConfig, PerSourceFilter\n", + "from autogen_agentchat.teams import (\n", + " DiGraphBuilder,\n", + " GraphFlow,\n", + ")\n", + "from autogen_agentchat.conditions import MaxMessageTermination\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\")\n", + "\n", + "# Agents\n", + "generator = AssistantAgent(\"generator\", model_client=model_client, system_message=\"Generate a list of creative ideas.\")\n", + "reviewer = AssistantAgent(\n", + " \"reviewer\",\n", + " model_client=model_client,\n", + " system_message=\"Review ideas and provide feedbacks, or just 'APPROVE' for final approval.\",\n", + ")\n", + "summarizer_core = AssistantAgent(\n", + " \"summary\", model_client=model_client, system_message=\"Summarize the user request and the final feedback.\"\n", + ")\n", + "\n", + "# Filtered summarizer\n", + "filtered_summarizer = MessageFilterAgent(\n", + " name=\"summary\",\n", + " wrapped_agent=summarizer_core,\n", + " filter=MessageFilterConfig(\n", + " per_source=[\n", + " PerSourceFilter(source=\"user\", position=\"first\", count=1),\n", + " PerSourceFilter(source=\"reviewer\", position=\"last\", count=1),\n", + " ]\n", + " ),\n", + ")\n", + "\n", + "# Build graph with conditional loop\n", + "builder = DiGraphBuilder()\n", + "builder.add_node(generator).add_node(reviewer).add_node(filtered_summarizer)\n", + "builder.add_edge(generator, reviewer)\n", + "builder.add_edge(reviewer, filtered_summarizer, condition=lambda msg: \"APPROVE\" in msg.to_model_text())\n", + "builder.add_edge(reviewer, generator, condition=lambda msg: \"APPROVE\" not in msg.to_model_text())\n", + "builder.set_entry_point(generator) # Set entry point to generator. Required if there are no source nodes.\n", + "graph = builder.build()\n", + "\n", + "termination_condition = MaxMessageTermination(10)\n", + "\n", + "# Create the flow\n", + "flow = GraphFlow(\n", + " participants=builder.get_participants(),\n", + " graph=graph,\n", + " termination_condition=termination_condition\n", + ")\n", + "\n", + "# Run the flow and pretty print the output in the console\n", + "await Console(flow.run_stream(task=\"Brainstorm ways to reduce plastic waste.\"))" + ] + }, + { + "cell_type": "markdown", + "id": "4b39f9d6", + "metadata": {}, + "source": [ + "## 🔁 Advanced Example: Cycles With Activation Group Examples\n", + "\n", + "The following examples demonstrate how to use `activation_group` and `activation_condition` to handle complex dependency patterns in cyclic graphs, especially when multiple paths lead to the same target node." + ] + }, + { + "cell_type": "markdown", + "id": "791a4c47", + "metadata": {}, + "source": [ + "### Example 1: Loop with Multiple Paths - \"All\" Activation (A→B→C→B)\n", + "\n", + "In this scenario, we have A → B → C → B, where B has two incoming edges (from A and from C). By default, B requires **all** its dependencies to be satisfied before executing.\n", + "\n", + "This example shows a review loop where both the initial input (A) and the feedback (C) must be processed before B can execute again." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "384f5831", + "metadata": {}, + "outputs": [], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.teams import DiGraphBuilder, GraphFlow\n", + "from autogen_agentchat.conditions import MaxMessageTermination\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "# Model client\n", + "client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\")\n", + "\n", + "# Create agents for A→B→C→B→E scenario\n", + "agent_a = AssistantAgent(\"A\", model_client=client, system_message=\"Start the process and provide initial input.\")\n", + "agent_b = AssistantAgent(\n", + " \"B\",\n", + " model_client=client,\n", + " system_message=\"Process input from A or feedback from C. Say 'CONTINUE' if it's from A or 'STOP' if it's from C.\",\n", + ")\n", + "agent_c = AssistantAgent(\"C\", model_client=client, system_message=\"Review B's output and provide feedback.\")\n", + "agent_e = AssistantAgent(\"E\", model_client=client, system_message=\"Finalize the process.\")\n", + "\n", + "# Build the graph with activation groups\n", + "builder = DiGraphBuilder()\n", + "builder.add_node(agent_a).add_node(agent_b).add_node(agent_c).add_node(agent_e)\n", + "\n", + "# A → B (initial path)\n", + "builder.add_edge(agent_a, agent_b, activation_group=\"initial\")\n", + "\n", + "# B → C\n", + "builder.add_edge(agent_b, agent_c, condition=\"CONTINUE\")\n", + "\n", + "# C → B (loop back - different activation group)\n", + "builder.add_edge(agent_c, agent_b, activation_group=\"feedback\")\n", + "\n", + "# B → E (exit condition)\n", + "builder.add_edge(agent_b, agent_e, condition=\"STOP\")\n", + "\n", + "termination_condition = MaxMessageTermination(10)\n", + "# Build and create flow\n", + "graph = builder.build()\n", + "flow = GraphFlow(participants=[agent_a, agent_b, agent_c, agent_e], graph=graph, termination_condition=termination_condition)\n", + "\n", + "print(\"=== Example 1: A→B→C→B with 'All' Activation ===\")\n", + "print(\"B will exit when it receives a message from C\")\n", + "# await Console(flow.run_stream(task=\"Start a review process for a document.\"))" + ] + }, + { + "cell_type": "markdown", + "id": "5dc08c64", + "metadata": {}, + "source": [ + "### Example 2: Loop with Multiple Paths - \"Any\" Activation (A→B→(C1,C2)→B)\n", + "\n", + "In this more complex scenario, we have A → B → (C1, C2) → B, where:\n", + "- B fans out to both C1 and C2 in parallel\n", + "- Both C1 and C2 feed back to B \n", + "- B uses \"any\" activation, meaning it executes as soon as **either** C1 or C2 completes\n", + "\n", + "This is useful for scenarios where you want the fastest response to trigger the next step.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "00f40293", + "metadata": {}, + "outputs": [], + "source": [ + "# Create agents for A→B→(C1,C2)→B scenario\n", + "agent_a2 = AssistantAgent(\"A\", model_client=client, system_message=\"Initiate a task that needs parallel processing.\")\n", + "agent_b2 = AssistantAgent(\n", + " \"B\",\n", + " model_client=client,\n", + " system_message=\"Coordinate parallel tasks. Say 'PROCESS' to start parallel work or 'DONE' to finish.\",\n", + ")\n", + "agent_c1 = AssistantAgent(\"C1\", model_client=client, system_message=\"Handle task type 1. Say 'C1_COMPLETE' when done.\")\n", + "agent_c2 = AssistantAgent(\"C2\", model_client=client, system_message=\"Handle task type 2. Say 'C2_COMPLETE' when done.\")\n", + "agent_e = AssistantAgent(\"E\", model_client=client, system_message=\"Finalize the process.\")\n", + "\n", + "# Build the graph with \"any\" activation\n", + "builder2 = DiGraphBuilder()\n", + "builder2.add_node(agent_a2).add_node(agent_b2).add_node(agent_c1).add_node(agent_c2).add_node(agent_e)\n", + "\n", + "# A → B (initial)\n", + "builder2.add_edge(agent_a2, agent_b2)\n", + "\n", + "# B → C1 and B → C2 (parallel fan-out)\n", + "builder2.add_edge(agent_b2, agent_c1, condition=\"PROCESS\")\n", + "builder2.add_edge(agent_b2, agent_c2, condition=\"PROCESS\")\n", + "\n", + "# B → E (exit condition)\n", + "builder2.add_edge(agent_b2, agent_e, condition=lambda msg: \"DONE\" in msg.to_model_text())\n", + "\n", + "# C1 → B and C2 → B (both in same activation group with \"any\" condition)\n", + "builder2.add_edge(\n", + " agent_c1, agent_b2, activation_group=\"loop_back_group\", activation_condition=\"any\", condition=\"C1_COMPLETE\"\n", + ")\n", + "\n", + "builder2.add_edge(\n", + " agent_c2, agent_b2, activation_group=\"loop_back_group\", activation_condition=\"any\", condition=\"C2_COMPLETE\"\n", + ")\n", + "\n", + "# Build and create flow\n", + "graph2 = builder2.build()\n", + "flow2 = GraphFlow(participants=[agent_a2, agent_b2, agent_c1, agent_c2, agent_e], graph=graph2)\n", + "\n", + "print(\"=== Example 2: A→B→(C1,C2)→B with 'Any' Activation ===\")\n", + "print(\"B will execute as soon as EITHER C1 OR C2 completes (whichever finishes first)\")\n", + "# await Console(flow2.run_stream(task=\"Start a parallel processing task.\"))" + ] + }, + { + "cell_type": "markdown", + "id": "7c56cd2e", + "metadata": {}, + "source": [ + "### Example 3: Mixed Activation Groups\n", + "\n", + "This example shows how different activation groups can coexist in the same graph. We have a scenario where:\n", + "- Node D receives inputs from multiple sources with different activation requirements\n", + "- Some dependencies use \"all\" activation (must wait for all inputs)\n", + "- Other dependencies use \"any\" activation (proceed on first input)\n", + "\n", + "This pattern is useful for complex workflows where different types of dependencies have different urgency levels.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "97f75ba1", + "metadata": {}, + "outputs": [], + "source": [ + "# Create agents for mixed activation scenario\n", + "agent_a3 = AssistantAgent(\"A\", model_client=client, system_message=\"Provide critical input that must be processed.\")\n", + "agent_b3 = AssistantAgent(\"B\", model_client=client, system_message=\"Provide secondary critical input.\")\n", + "agent_c3 = AssistantAgent(\"C\", model_client=client, system_message=\"Provide optional quick input.\")\n", + "agent_d3 = AssistantAgent(\"D\", model_client=client, system_message=\"Process inputs based on different priority levels.\")\n", + "\n", + "# Build graph with mixed activation groups\n", + "builder3 = DiGraphBuilder()\n", + "builder3.add_node(agent_a3).add_node(agent_b3).add_node(agent_c3).add_node(agent_d3)\n", + "\n", + "# Critical inputs that must ALL be present (activation_group=\"critical\", activation_condition=\"all\")\n", + "builder3.add_edge(agent_a3, agent_d3, activation_group=\"critical\", activation_condition=\"all\")\n", + "builder3.add_edge(agent_b3, agent_d3, activation_group=\"critical\", activation_condition=\"all\")\n", + "\n", + "# Optional input that can trigger execution on its own (activation_group=\"optional\", activation_condition=\"any\")\n", + "builder3.add_edge(agent_c3, agent_d3, activation_group=\"optional\", activation_condition=\"any\")\n", + "\n", + "# Build and create flow\n", + "graph3 = builder3.build()\n", + "flow3 = GraphFlow(participants=[agent_a3, agent_b3, agent_c3, agent_d3], graph=graph3)\n", + "\n", + "print(\"=== Example 3: Mixed Activation Groups ===\")\n", + "print(\"D will execute when:\")\n", + "print(\"- BOTH A AND B complete (critical group with 'all' activation), OR\")\n", + "print(\"- C completes (optional group with 'any' activation)\")\n", + "print(\"This allows for both required dependencies and fast-path triggers.\")\n", + "# await Console(flow3.run_stream(task=\"Process inputs with mixed priority levels.\"))" + ] + }, + { + "cell_type": "markdown", + "id": "e329fe57", + "metadata": {}, + "source": [ + "### Key Takeaways for Activation Groups\n", + "\n", + "1. **`activation_group`**: Groups edges that point to the same target node, allowing you to define different dependency patterns.\n", + "\n", + "2. **`activation_condition`**: \n", + " - `\"all\"` (default): Target node waits for ALL edges in the group to be satisfied\n", + " - `\"any\"`: Target node executes as soon as ANY edge in the group is satisfied\n", + "\n", + "3. **Use Cases**:\n", + " - **Cycles with multiple entry points**: Different activation groups prevent conflicts\n", + " - **Priority-based execution**: Mix \"all\" and \"any\" conditions for different urgency levels \n", + " - **Parallel processing with early termination**: Use \"any\" to proceed with the fastest result\n", + "\n", + "4. **Best Practices**:\n", + " - Use descriptive group names (`\"critical\"`, `\"optional\"`, `\"feedback\"`, etc.)\n", + " - Keep activation conditions consistent within the same group\n", + " - Test your graph logic with different execution paths\n", + "\n", + "These patterns enable sophisticated workflow control while maintaining clear, understandable execution semantics." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "python", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.7" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/index.md b/python/docs/src/user-guide/agentchat-user-guide/index.md similarity index 93% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/index.md rename to python/docs/src/user-guide/agentchat-user-guide/index.md index fc84f2206a48..016111df7a30 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/index.md +++ b/python/docs/src/user-guide/agentchat-user-guide/index.md @@ -68,6 +68,13 @@ Multi-agent coordination through a shared context and localized, tool-based sele Get started with Magentic-One ::: +:::{grid-item-card} {fas}`sitemap;pst-color-primary` GraphFlow (Workflow) +:link: ./graph-flow.html +:link-alt: GraphFlow: Multi-agent workflows through a directed graph of agents. + +Multi-agent workflows through a directed graph of agents. +::: + :::{grid-item-card} {fas}`brain;pst-color-primary` Memory :link: ./memory.html :link-alt: Memory: Add memory capabilities to your agents @@ -138,6 +145,7 @@ custom-agents selector-group-chat swarm magentic-one +graph-flow memory logging serialize-components diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/installation.md b/python/docs/src/user-guide/agentchat-user-guide/installation.md similarity index 88% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/installation.md rename to python/docs/src/user-guide/agentchat-user-guide/installation.md index 8171ff16ac77..c84e45d5b2da 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/installation.md +++ b/python/docs/src/user-guide/agentchat-user-guide/installation.md @@ -17,13 +17,19 @@ When installing AgentChat locally, we recommend using a virtual environment for Create and activate: +Linux/Mac: ```bash -# On Windows, change `python3` to `python` (if `python` is Python 3). python3 -m venv .venv -# On Windows, change `bin` to `scripts`. source .venv/bin/activate ``` +Windows command-line: +```batch +# The command may be `python3` instead of `python` depending on your setup +python -m venv .venv +.venv\Scripts\activate.bat +``` + To deactivate later, run: ```bash diff --git a/python/docs/src/user-guide/agentchat-user-guide/jaeger.png b/python/docs/src/user-guide/agentchat-user-guide/jaeger.png new file mode 100644 index 000000000000..cbda69c199ee --- /dev/null +++ b/python/docs/src/user-guide/agentchat-user-guide/jaeger.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d0e3f50140f8bef32e2f626792490b2c1c4728e5ec10130241cf3f49ada44a20 +size 167573 diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/logging.md b/python/docs/src/user-guide/agentchat-user-guide/logging.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/logging.md rename to python/docs/src/user-guide/agentchat-user-guide/logging.md diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/magentic-one.md b/python/docs/src/user-guide/agentchat-user-guide/magentic-one.md similarity index 94% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/magentic-one.md rename to python/docs/src/user-guide/agentchat-user-guide/magentic-one.md index 46b90ad453f7..6de052d1a9b0 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/magentic-one.md +++ b/python/docs/src/user-guide/agentchat-user-guide/magentic-one.md @@ -121,11 +121,23 @@ import asyncio from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.teams.magentic_one import MagenticOne from autogen_agentchat.ui import Console +from autogen_agentchat.agents import ApprovalRequest, ApprovalResponse + + +def approval_func(request: ApprovalRequest) -> ApprovalResponse: + """Simple approval function that requests user input before code execution.""" + print(f"Code to execute:\n{request.code}") + user_input = input("Do you approve this code execution? (y/n): ").strip().lower() + if user_input == 'y': + return ApprovalResponse(approved=True, reason="User approved the code execution") + else: + return ApprovalResponse(approved=False, reason="User denied the code execution") async def example_usage(): client = OpenAIChatCompletionClient(model="gpt-4o") - m1 = MagenticOne(client=client) + # Enable code execution approval for security + m1 = MagenticOne(client=client, approval_func=approval_func) task = "Write a Python script to fetch data from an API." result = await Console(m1.run_stream(task=task)) print(result) diff --git a/python/docs/src/user-guide/agentchat-user-guide/memory.ipynb b/python/docs/src/user-guide/agentchat-user-guide/memory.ipynb new file mode 100644 index 000000000000..35b14939e82b --- /dev/null +++ b/python/docs/src/user-guide/agentchat-user-guide/memory.ipynb @@ -0,0 +1,756 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Memory and RAG\n", + "\n", + "There are several use cases where it is valuable to maintain a _store_ of useful facts that can be intelligently added to the context of the agent just before a specific step. The typically use case here is a RAG pattern where a query is used to retrieve relevant information from a database that is then added to the agent's context.\n", + "\n", + "\n", + "AgentChat provides a {py:class}`~autogen_core.memory.Memory` protocol that can be extended to provide this functionality. The key methods are `query`, `update_context`, `add`, `clear`, and `close`. \n", + "\n", + "- `add`: add new entries to the memory store\n", + "- `query`: retrieve relevant information from the memory store \n", + "- `update_context`: mutate an agent's internal `model_context` by adding the retrieved information (used in the {py:class}`~autogen_agentchat.agents.AssistantAgent` class) \n", + "- `clear`: clear all entries from the memory store\n", + "- `close`: clean up any resources used by the memory store \n", + "\n", + "\n", + "## ListMemory Example\n", + "\n", + "{py:class}`~autogen_core.memory.ListMemory` is provided as an example implementation of the {py:class}`~autogen_core.memory.Memory` protocol. It is a simple list-based memory implementation that maintains memories in chronological order, appending the most recent memories to the model's context. The implementation is designed to be straightforward and predictable, making it easy to understand and debug.\n", + "In the following example, we will use ListMemory to maintain a memory bank of user preferences and demonstrate how it can be used to provide consistent context for agent responses over time." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_core.memory import ListMemory, MemoryContent, MemoryMimeType\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "# Initialize user memory\n", + "user_memory = ListMemory()\n", + "\n", + "# Add user preferences to memory\n", + "await user_memory.add(MemoryContent(content=\"The weather should be in metric units\", mime_type=MemoryMimeType.TEXT))\n", + "\n", + "await user_memory.add(MemoryContent(content=\"Meal recipe must be vegan\", mime_type=MemoryMimeType.TEXT))\n", + "\n", + "\n", + "async def get_weather(city: str, units: str = \"imperial\") -> str:\n", + " if units == \"imperial\":\n", + " return f\"The weather in {city} is 73 °F and Sunny.\"\n", + " elif units == \"metric\":\n", + " return f\"The weather in {city} is 23 °C and Sunny.\"\n", + " else:\n", + " return f\"Sorry, I don't know the weather in {city}.\"\n", + "\n", + "\n", + "assistant_agent = AssistantAgent(\n", + " name=\"assistant_agent\",\n", + " model_client=OpenAIChatCompletionClient(\n", + " model=\"gpt-4o-2024-08-06\",\n", + " ),\n", + " tools=[get_weather],\n", + " memory=[user_memory],\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- TextMessage (user) ----------\n", + "What is the weather in New York?\n", + "---------- MemoryQueryEvent (assistant_agent) ----------\n", + "[MemoryContent(content='The weather should be in metric units', mime_type=, metadata=None), MemoryContent(content='Meal recipe must be vegan', mime_type=, metadata=None)]\n", + "---------- ToolCallRequestEvent (assistant_agent) ----------\n", + "[FunctionCall(id='call_33uMqZO6hwOfEpJavP9GW9LI', arguments='{\"city\":\"New York\",\"units\":\"metric\"}', name='get_weather')]\n", + "---------- ToolCallExecutionEvent (assistant_agent) ----------\n", + "[FunctionExecutionResult(content='The weather in New York is 23 °C and Sunny.', name='get_weather', call_id='call_33uMqZO6hwOfEpJavP9GW9LI', is_error=False)]\n", + "---------- ToolCallSummaryMessage (assistant_agent) ----------\n", + "The weather in New York is 23 °C and Sunny.\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, created_at=datetime.datetime(2025, 7, 1, 23, 53, 8, 867845, tzinfo=datetime.timezone.utc), content='What is the weather in New York?', type='TextMessage'), MemoryQueryEvent(source='assistant_agent', models_usage=None, metadata={}, created_at=datetime.datetime(2025, 7, 1, 23, 53, 8, 869589, tzinfo=datetime.timezone.utc), content=[MemoryContent(content='The weather should be in metric units', mime_type=, metadata=None), MemoryContent(content='Meal recipe must be vegan', mime_type=, metadata=None)], type='MemoryQueryEvent'), ToolCallRequestEvent(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=123, completion_tokens=19), metadata={}, created_at=datetime.datetime(2025, 7, 1, 23, 53, 10, 240626, tzinfo=datetime.timezone.utc), content=[FunctionCall(id='call_33uMqZO6hwOfEpJavP9GW9LI', arguments='{\"city\":\"New York\",\"units\":\"metric\"}', name='get_weather')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='assistant_agent', models_usage=None, metadata={}, created_at=datetime.datetime(2025, 7, 1, 23, 53, 10, 242633, tzinfo=datetime.timezone.utc), content=[FunctionExecutionResult(content='The weather in New York is 23 °C and Sunny.', name='get_weather', call_id='call_33uMqZO6hwOfEpJavP9GW9LI', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='assistant_agent', models_usage=None, metadata={}, created_at=datetime.datetime(2025, 7, 1, 23, 53, 10, 243722, tzinfo=datetime.timezone.utc), content='The weather in New York is 23 °C and Sunny.', type='ToolCallSummaryMessage')], stop_reason=None)" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Run the agent with a task.\n", + "stream = assistant_agent.run_stream(task=\"What is the weather in New York?\")\n", + "await Console(stream)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can inspect that the `assistant_agent` model_context is actually updated with the retrieved memory entries. The `transform` method is used to format the retrieved memory entries into a string that can be used by the agent. In this case, we simply concatenate the content of each memory entry into a single string." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[UserMessage(content='What is the weather in New York?', source='user', type='UserMessage'),\n", + " SystemMessage(content='\\nRelevant memory content (in chronological order):\\n1. The weather should be in metric units\\n2. Meal recipe must be vegan\\n', type='SystemMessage'),\n", + " AssistantMessage(content=[FunctionCall(id='call_33uMqZO6hwOfEpJavP9GW9LI', arguments='{\"city\":\"New York\",\"units\":\"metric\"}', name='get_weather')], thought=None, source='assistant_agent', type='AssistantMessage'),\n", + " FunctionExecutionResultMessage(content=[FunctionExecutionResult(content='The weather in New York is 23 °C and Sunny.', name='get_weather', call_id='call_33uMqZO6hwOfEpJavP9GW9LI', is_error=False)], type='FunctionExecutionResultMessage')]" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "await assistant_agent._model_context.get_messages()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We see above that the weather is returned in Centigrade as stated in the user preferences. \n", + "\n", + "Similarly, assuming we ask a separate question about generating a meal plan, the agent is able to retrieve relevant information from the memory store and provide a personalized (vegan) response." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- TextMessage (user) ----------\n", + "Write brief meal recipe with broth\n", + "---------- MemoryQueryEvent (assistant_agent) ----------\n", + "[MemoryContent(content='The weather should be in metric units', mime_type=, metadata=None), MemoryContent(content='Meal recipe must be vegan', mime_type=, metadata=None)]\n", + "---------- TextMessage (assistant_agent) ----------\n", + "Here's a brief vegan meal recipe using broth:\n", + "\n", + "**Vegan Vegetable Broth Soup**\n", + "\n", + "**Ingredients:**\n", + "- 1 tablespoon olive oil\n", + "- 1 onion, chopped\n", + "- 3 cloves garlic, minced\n", + "- 2 carrots, sliced\n", + "- 2 celery stalks, sliced\n", + "- 1 zucchini, chopped\n", + "- 1 cup mushrooms, sliced\n", + "- 1 cup kale or spinach, chopped\n", + "- 1 can (400g) diced tomatoes\n", + "- 4 cups vegetable broth\n", + "- 1 teaspoon dried thyme\n", + "- Salt and pepper to taste\n", + "- Fresh parsley, chopped (for garnish)\n", + "\n", + "**Instructions:**\n", + "1. Heat olive oil in a large pot over medium heat. Add the onion and garlic, and sauté until soft.\n", + "2. Add the carrots, celery, zucchini, and mushrooms. Cook for about 5 minutes until the vegetables begin to soften.\n", + "3. Add the diced tomatoes, vegetable broth, and dried thyme. Bring to a boil.\n", + "4. Reduce heat and let it simmer for about 20 minutes, or until the vegetables are tender.\n", + "5. Stir in the chopped kale or spinach and cook for another 5 minutes.\n", + "6. Season with salt and pepper to taste.\n", + "7. Serve hot, garnished with fresh parsley.\n", + "\n", + "Enjoy your comforting vegan vegetable broth soup!\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, created_at=datetime.datetime(2025, 7, 1, 23, 53, 10, 256897, tzinfo=datetime.timezone.utc), content='Write brief meal recipe with broth', type='TextMessage'), MemoryQueryEvent(source='assistant_agent', models_usage=None, metadata={}, created_at=datetime.datetime(2025, 7, 1, 23, 53, 10, 258468, tzinfo=datetime.timezone.utc), content=[MemoryContent(content='The weather should be in metric units', mime_type=, metadata=None), MemoryContent(content='Meal recipe must be vegan', mime_type=, metadata=None)], type='MemoryQueryEvent'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=205, completion_tokens=266), metadata={}, created_at=datetime.datetime(2025, 7, 1, 23, 53, 14, 67151, tzinfo=datetime.timezone.utc), content=\"Here's a brief vegan meal recipe using broth:\\n\\n**Vegan Vegetable Broth Soup**\\n\\n**Ingredients:**\\n- 1 tablespoon olive oil\\n- 1 onion, chopped\\n- 3 cloves garlic, minced\\n- 2 carrots, sliced\\n- 2 celery stalks, sliced\\n- 1 zucchini, chopped\\n- 1 cup mushrooms, sliced\\n- 1 cup kale or spinach, chopped\\n- 1 can (400g) diced tomatoes\\n- 4 cups vegetable broth\\n- 1 teaspoon dried thyme\\n- Salt and pepper to taste\\n- Fresh parsley, chopped (for garnish)\\n\\n**Instructions:**\\n1. Heat olive oil in a large pot over medium heat. Add the onion and garlic, and sauté until soft.\\n2. Add the carrots, celery, zucchini, and mushrooms. Cook for about 5 minutes until the vegetables begin to soften.\\n3. Add the diced tomatoes, vegetable broth, and dried thyme. Bring to a boil.\\n4. Reduce heat and let it simmer for about 20 minutes, or until the vegetables are tender.\\n5. Stir in the chopped kale or spinach and cook for another 5 minutes.\\n6. Season with salt and pepper to taste.\\n7. Serve hot, garnished with fresh parsley.\\n\\nEnjoy your comforting vegan vegetable broth soup!\", type='TextMessage')], stop_reason=None)" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "stream = assistant_agent.run_stream(task=\"Write brief meal recipe with broth\")\n", + "await Console(stream)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Custom Memory Stores (Vector DBs, etc.)\n", + "\n", + "You can build on the `Memory` protocol to implement more complex memory stores. For example, you could implement a custom memory store that uses a vector database to store and retrieve information, or a memory store that uses a machine learning model to generate personalized responses based on the user's preferences etc.\n", + "\n", + "Specifically, you will need to overload the `add`, `query` and `update_context` methods to implement the desired functionality and pass the memory store to your agent.\n", + "\n", + "\n", + "Currently the following example memory stores are available as part of the {py:class}`~autogen_ext` extensions package.\n", + "\n", + "- `autogen_ext.memory.chromadb.ChromaDBVectorMemory`: A memory store that uses a vector database to store and retrieve information.\n", + "\n", + "- `autogen_ext.memory.chromadb.SentenceTransformerEmbeddingFunctionConfig`: A configuration class for the SentenceTransformer embedding function used by the `ChromaDBVectorMemory` store. Note that other embedding functions such as `autogen_ext.memory.openai.OpenAIEmbeddingFunctionConfig` can also be used with the `ChromaDBVectorMemory` store.\n", + "\n", + "- `autogen_ext.memory.redis.RedisMemory`: A memory store that uses a Redis vector database to store and retrieve information.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- TextMessage (user) ----------\n", + "What is the weather in New York?\n", + "---------- MemoryQueryEvent (assistant_agent) ----------\n", + "[MemoryContent(content='The weather should be in metric units', mime_type='MemoryMimeType.TEXT', metadata={'category': 'preferences', 'mime_type': 'MemoryMimeType.TEXT', 'type': 'units', 'score': 0.4342913031578064, 'id': 'b8a70e90-a39f-47ed-ab7b-5a274009d9f0'}), MemoryContent(content='The weather should be in metric units', mime_type='MemoryMimeType.TEXT', metadata={'mime_type': 'MemoryMimeType.TEXT', 'type': 'units', 'category': 'preferences', 'score': 0.4342913031578064, 'id': 'b240f12a-1440-42d1-8f5e-3d8a388363f2'})]\n", + "---------- ToolCallRequestEvent (assistant_agent) ----------\n", + "[FunctionCall(id='call_YmKqq1nWXgAkAAyXWWk9YpFW', arguments='{\"city\":\"New York\",\"units\":\"metric\"}', name='get_weather')]\n", + "---------- ToolCallExecutionEvent (assistant_agent) ----------\n", + "[FunctionExecutionResult(content='The weather in New York is 23 °C and Sunny.', name='get_weather', call_id='call_YmKqq1nWXgAkAAyXWWk9YpFW', is_error=False)]\n", + "---------- ToolCallSummaryMessage (assistant_agent) ----------\n", + "The weather in New York is 23 °C and Sunny.\n" + ] + } + ], + "source": [ + "import tempfile\n", + "\n", + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_core.memory import MemoryContent, MemoryMimeType\n", + "from autogen_ext.memory.chromadb import (\n", + " ChromaDBVectorMemory,\n", + " PersistentChromaDBVectorMemoryConfig,\n", + " SentenceTransformerEmbeddingFunctionConfig,\n", + ")\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "# Use a temporary directory for ChromaDB persistence\n", + "with tempfile.TemporaryDirectory() as tmpdir:\n", + " chroma_user_memory = ChromaDBVectorMemory(\n", + " config=PersistentChromaDBVectorMemoryConfig(\n", + " collection_name=\"preferences\",\n", + " persistence_path=tmpdir, # Use the temp directory here\n", + " k=2, # Return top k results\n", + " score_threshold=0.4, # Minimum similarity score\n", + " embedding_function_config=SentenceTransformerEmbeddingFunctionConfig(\n", + " model_name=\"all-MiniLM-L6-v2\" # Use default model for testing\n", + " ),\n", + " )\n", + " )\n", + " # Add user preferences to memory\n", + " await chroma_user_memory.add(\n", + " MemoryContent(\n", + " content=\"The weather should be in metric units\",\n", + " mime_type=MemoryMimeType.TEXT,\n", + " metadata={\"category\": \"preferences\", \"type\": \"units\"},\n", + " )\n", + " )\n", + "\n", + " await chroma_user_memory.add(\n", + " MemoryContent(\n", + " content=\"Meal recipe must be vegan\",\n", + " mime_type=MemoryMimeType.TEXT,\n", + " metadata={\"category\": \"preferences\", \"type\": \"dietary\"},\n", + " )\n", + " )\n", + "\n", + " model_client = OpenAIChatCompletionClient(\n", + " model=\"gpt-4o\",\n", + " )\n", + "\n", + " # Create assistant agent with ChromaDB memory\n", + " assistant_agent = AssistantAgent(\n", + " name=\"assistant_agent\",\n", + " model_client=model_client,\n", + " tools=[get_weather],\n", + " memory=[chroma_user_memory],\n", + " )\n", + "\n", + " stream = assistant_agent.run_stream(task=\"What is the weather in New York?\")\n", + " await Console(stream)\n", + "\n", + " await model_client.close()\n", + " await chroma_user_memory.close()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Note that you can also serialize the ChromaDBVectorMemory and save it to disk." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'{\"provider\":\"autogen_ext.memory.chromadb.ChromaDBVectorMemory\",\"component_type\":\"memory\",\"version\":1,\"component_version\":1,\"description\":\"Store and retrieve memory using vector similarity search powered by ChromaDB.\",\"label\":\"ChromaDBVectorMemory\",\"config\":{\"client_type\":\"persistent\",\"collection_name\":\"preferences\",\"distance_metric\":\"cosine\",\"k\":2,\"score_threshold\":0.4,\"allow_reset\":false,\"tenant\":\"default_tenant\",\"database\":\"default_database\",\"persistence_path\":\"/Users/justin.cechmanek/.chromadb_autogen\"}}'" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "chroma_user_memory.dump_component().model_dump_json()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Redis Memory\n", + "You can perform the same persistent memory storage using Redis. Note, you will need to have a running Redis instance to connect to.\n", + "\n", + "See {py:class}`~autogen_ext.memory.redis.RedisMemory` for instructions to run Redis locally or via Docker." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- TextMessage (user) ----------\n", + "What is the weather in New York?\n", + "---------- MemoryQueryEvent (assistant_agent) ----------\n", + "[MemoryContent(content='The weather should be in metric units', mime_type=, metadata={'category': 'preferences', 'type': 'units'})]\n", + "---------- ToolCallRequestEvent (assistant_agent) ----------\n", + "[FunctionCall(id='call_1R6wV3uDOK8mGK2Vh2t0h4ld', arguments='{\"city\":\"New York\",\"units\":\"metric\"}', name='get_weather')]\n", + "---------- ToolCallExecutionEvent (assistant_agent) ----------\n", + "[FunctionExecutionResult(content='The weather in New York is 23 °C and Sunny.', name='get_weather', call_id='call_1R6wV3uDOK8mGK2Vh2t0h4ld', is_error=False)]\n", + "---------- ToolCallSummaryMessage (assistant_agent) ----------\n", + "The weather in New York is 23 °C and Sunny.\n" + ] + } + ], + "source": [ + "from logging import WARNING, getLogger\n", + "\n", + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_core.memory import MemoryContent, MemoryMimeType\n", + "from autogen_ext.memory.redis import RedisMemory, RedisMemoryConfig\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "logger = getLogger()\n", + "logger.setLevel(WARNING)\n", + "\n", + "# Initailize Redis memory\n", + "redis_memory = RedisMemory(\n", + " config=RedisMemoryConfig(\n", + " redis_url=\"redis://localhost:6379\",\n", + " index_name=\"chat_history\",\n", + " prefix=\"memory\",\n", + " )\n", + ")\n", + "\n", + "# Add user preferences to memory\n", + "await redis_memory.add(\n", + " MemoryContent(\n", + " content=\"The weather should be in metric units\",\n", + " mime_type=MemoryMimeType.TEXT,\n", + " metadata={\"category\": \"preferences\", \"type\": \"units\"},\n", + " )\n", + ")\n", + "\n", + "await redis_memory.add(\n", + " MemoryContent(\n", + " content=\"Meal recipe must be vegan\",\n", + " mime_type=MemoryMimeType.TEXT,\n", + " metadata={\"category\": \"preferences\", \"type\": \"dietary\"},\n", + " )\n", + ")\n", + "\n", + "model_client = OpenAIChatCompletionClient(\n", + " model=\"gpt-4o\",\n", + ")\n", + "\n", + "# Create assistant agent with ChromaDB memory\n", + "assistant_agent = AssistantAgent(\n", + " name=\"assistant_agent\",\n", + " model_client=model_client,\n", + " tools=[get_weather],\n", + " memory=[redis_memory],\n", + ")\n", + "\n", + "stream = assistant_agent.run_stream(task=\"What is the weather in New York?\")\n", + "await Console(stream)\n", + "\n", + "await model_client.close()\n", + "await redis_memory.close()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## RAG Agent: Putting It All Together\n", + "\n", + "The RAG (Retrieval Augmented Generation) pattern which is common in building AI systems encompasses two distinct phases:\n", + "\n", + "1. **Indexing**: Loading documents, chunking them, and storing them in a vector database\n", + "2. **Retrieval**: Finding and using relevant chunks during conversation runtime\n", + "\n", + "In our previous examples, we manually added items to memory and passed them to our agents. In practice, the indexing process is usually automated and based on much larger document sources like product documentation, internal files, or knowledge bases.\n", + "\n", + "> Note: The quality of a RAG system is dependent on the quality of the chunking and retrieval process (models, embeddings, etc.). You may need to experiement with more advanced chunking and retrieval models to get the best results.\n", + "\n", + "### Building a Simple RAG Agent\n", + "\n", + "To begin, let's create a simple document indexer that we will used to load documents, chunk them, and store them in a `ChromaDBVectorMemory` memory store. " + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [], + "source": [ + "import re\n", + "from typing import List\n", + "\n", + "import aiofiles\n", + "import aiohttp\n", + "from autogen_core.memory import Memory, MemoryContent, MemoryMimeType\n", + "\n", + "\n", + "class SimpleDocumentIndexer:\n", + " \"\"\"Basic document indexer for AutoGen Memory.\"\"\"\n", + "\n", + " def __init__(self, memory: Memory, chunk_size: int = 1500) -> None:\n", + " self.memory = memory\n", + " self.chunk_size = chunk_size\n", + "\n", + " async def _fetch_content(self, source: str) -> str:\n", + " \"\"\"Fetch content from URL or file.\"\"\"\n", + " if source.startswith((\"http://\", \"https://\")):\n", + " async with aiohttp.ClientSession() as session:\n", + " async with session.get(source) as response:\n", + " return await response.text()\n", + " else:\n", + " async with aiofiles.open(source, \"r\", encoding=\"utf-8\") as f:\n", + " return await f.read()\n", + "\n", + " def _strip_html(self, text: str) -> str:\n", + " \"\"\"Remove HTML tags and normalize whitespace.\"\"\"\n", + " text = re.sub(r\"<[^>]*>\", \" \", text)\n", + " text = re.sub(r\"\\s+\", \" \", text)\n", + " return text.strip()\n", + "\n", + " def _split_text(self, text: str) -> List[str]:\n", + " \"\"\"Split text into fixed-size chunks.\"\"\"\n", + " chunks: list[str] = []\n", + " # Just split text into fixed-size chunks\n", + " for i in range(0, len(text), self.chunk_size):\n", + " chunk = text[i : i + self.chunk_size]\n", + " chunks.append(chunk.strip())\n", + " return chunks\n", + "\n", + " async def index_documents(self, sources: List[str]) -> int:\n", + " \"\"\"Index documents into memory.\"\"\"\n", + " total_chunks = 0\n", + "\n", + " for source in sources:\n", + " try:\n", + " content = await self._fetch_content(source)\n", + "\n", + " # Strip HTML if content appears to be HTML\n", + " if \"<\" in content and \">\" in content:\n", + " content = self._strip_html(content)\n", + "\n", + " chunks = self._split_text(content)\n", + "\n", + " for i, chunk in enumerate(chunks):\n", + " await self.memory.add(\n", + " MemoryContent(\n", + " content=chunk, mime_type=MemoryMimeType.TEXT, metadata={\"source\": source, \"chunk_index\": i}\n", + " )\n", + " )\n", + "\n", + " total_chunks += len(chunks)\n", + "\n", + " except Exception as e:\n", + " print(f\"Error indexing {source}: {str(e)}\")\n", + "\n", + " return total_chunks" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + " \n", + "Now let's use our indexer with ChromaDBVectorMemory to build a complete RAG agent:\n" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Indexed 70 chunks from 4 AutoGen documents\n" + ] + } + ], + "source": [ + "import os\n", + "from pathlib import Path\n", + "\n", + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_ext.memory.chromadb import ChromaDBVectorMemory, PersistentChromaDBVectorMemoryConfig\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "# Initialize vector memory\n", + "\n", + "rag_memory = ChromaDBVectorMemory(\n", + " config=PersistentChromaDBVectorMemoryConfig(\n", + " collection_name=\"autogen_docs\",\n", + " persistence_path=os.path.join(str(Path.home()), \".chromadb_autogen\"),\n", + " k=3, # Return top 3 results\n", + " score_threshold=0.4, # Minimum similarity score\n", + " )\n", + ")\n", + "\n", + "await rag_memory.clear() # Clear existing memory\n", + "\n", + "\n", + "# Index AutoGen documentation\n", + "async def index_autogen_docs() -> None:\n", + " indexer = SimpleDocumentIndexer(memory=rag_memory)\n", + " sources = [\n", + " \"https://raw.githubusercontent.com/microsoft/autogen/main/README.md\",\n", + " \"https://microsoft.github.io/autogen/dev/user-guide/agentchat-user-guide/tutorial/agents.html\",\n", + " \"https://microsoft.github.io/autogen/dev/user-guide/agentchat-user-guide/tutorial/teams.html\",\n", + " \"https://microsoft.github.io/autogen/dev/user-guide/agentchat-user-guide/tutorial/termination.html\",\n", + " ]\n", + " chunks: int = await indexer.index_documents(sources)\n", + " print(f\"Indexed {chunks} chunks from {len(sources)} AutoGen documents\")\n", + "\n", + "\n", + "await index_autogen_docs()" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- TextMessage (user) ----------\n", + "What is AgentChat?\n", + "---------- MemoryQueryEvent (rag_assistant) ----------\n", + "[MemoryContent(content='e of the AssistantAgent , we can now proceed to the next section to learn about the teams feature in AgentChat. previous Messages next Teams On this page Assistant Agent Getting Result Multi-Modal Input Streaming Messages Using Tools and Workbench Built-in Tools and Workbench Function Tool Model Context Protocol (MCP) Workbench Agent as a Tool Parallel Tool Calls Tool Iterations Structured Output Streaming Tokens Using Model Context Other Preset Agents Next Step Edit on GitHub Show Source so the DOM is not blocked --> © Copyright 2024, Microsoft. Privacy Policy | Consumer Health Privacy Built with the PyData Sphinx Theme 0.16.0.', mime_type='MemoryMimeType.TEXT', metadata={'chunk_index': 16, 'mime_type': 'MemoryMimeType.TEXT', 'source': 'https://microsoft.github.io/autogen/dev/user-guide/agentchat-user-guide/tutorial/agents.html', 'score': 0.6237251460552216, 'id': '6457da13-1c25-44f0-bea3-158e5c0c5bb4'}), MemoryContent(content='h Literature Review API Reference PyPi Source AgentChat Agents Agents # AutoGen AgentChat provides a set of preset Agents, each with variations in how an agent might respond to messages. All agents share the following attributes and methods: name : The unique name of the agent. description : The description of the agent in text. run : The method that runs the agent given a task as a string or a list of messages, and returns a TaskResult . Agents are expected to be stateful and this method is expected to be called with new messages, not complete history . run_stream : Same as run() but returns an iterator of messages that subclass BaseAgentEvent or BaseChatMessage followed by a TaskResult as the last item. See autogen_agentchat.messages for more information on AgentChat message types. Assistant Agent # AssistantAgent is a built-in agent that uses a language model and has the ability to use tools. Warning AssistantAgent is a “kitchen sink” agent for prototyping and educational purpose – it is very general. Make sure you read the documentation and implementation to understand the design choices. Once you fully understand the design, you may want to implement your own agent. See Custom Agent . from autogen_agentchat.agents import AssistantAgent from autogen_agentchat.messages import StructuredMessage from autogen_agentchat.ui import Console from autogen_ext.models.openai import OpenAIChatCompletionClient # Define a tool that searches the web for information. # For simplicity, we', mime_type='MemoryMimeType.TEXT', metadata={'chunk_index': 1, 'mime_type': 'MemoryMimeType.TEXT', 'source': 'https://microsoft.github.io/autogen/dev/user-guide/agentchat-user-guide/tutorial/agents.html', 'score': 0.6212755441665649, 'id': 'ab3a553f-bb69-41ff-b6a9-8397b4cb3cb1'}), MemoryContent(content='Literature Review API Reference PyPi Source AgentChat Teams Teams # In this section you’ll learn how to create a multi-agent team (or simply team) using AutoGen. A team is a group of agents that work together to achieve a common goal. We’ll first show you how to create and run a team. We’ll then explain how to observe the team’s behavior, which is crucial for debugging and understanding the team’s performance, and common operations to control the team’s behavior. AgentChat supports several team presets: RoundRobinGroupChat : A team that runs a group chat with participants taking turns in a round-robin fashion (covered on this page). Tutorial SelectorGroupChat : A team that selects the next speaker using a ChatCompletion model after each message. Tutorial MagenticOneGroupChat : A generalist multi-agent system for solving open-ended web and file-based tasks across a variety of domains. Tutorial Swarm : A team that uses HandoffMessage to signal transitions between agents. Tutorial Note When should you use a team? Teams are for complex tasks that require collaboration and diverse expertise. However, they also demand more scaffolding to steer compared to single agents. While AutoGen simplifies the process of working with teams, start with a single agent for simpler tasks, and transition to a multi-agent team when a single agent proves inadequate. Ensure that you have optimized your single agent with the appropriate tools and instructions before moving to a team-based approach. Cre', mime_type='MemoryMimeType.TEXT', metadata={'mime_type': 'MemoryMimeType.TEXT', 'chunk_index': 1, 'source': 'https://microsoft.github.io/autogen/dev/user-guide/agentchat-user-guide/tutorial/teams.html', 'score': 0.5267025232315063, 'id': '554b20a9-e041-4ac6-b2f1-11261336861c'})]\n", + "---------- TextMessage (rag_assistant) ----------\n", + "AgentChat is a framework that provides a set of preset agents designed to handle conversations and tasks using a variety of response strategies. It includes features for managing individual agents as well as creating teams of agents that can work collaboratively on complex goals. These agents are stateful, meaning they can manage and track ongoing conversations. AgentChat also includes agents that can utilize tools to enhance their capabilities.\n", + "\n", + "Key features of AgentChat include:\n", + "- **Preset Agents**: These agents are pre-configured with specific behavior patterns for handling tasks and messages.\n", + "- **Agent Attributes and Methods**: Each agent has a unique name and description, and methods like `run` and `run_stream` to execute tasks and handle messages.\n", + "- **AssistantAgent**: A built-in general-purpose agent used primarily for prototyping and educational purposes.\n", + "- **Team Configurations**: AgentChat allows for the creation of multi-agent teams for tasks that are too complex for a single agent. Teams run in preset formats like RoundRobinGroupChat or Swarm, providing structured interaction among agents.\n", + "\n", + "Overall, AgentChat is designed for flexible deployment of conversational agents, either singly or in groups, across a variety of tasks. \n", + "\n", + "TERMINATE\n" + ] + } + ], + "source": [ + "# Create our RAG assistant agent\n", + "rag_assistant = AssistantAgent(\n", + " name=\"rag_assistant\", model_client=OpenAIChatCompletionClient(model=\"gpt-4o\"), memory=[rag_memory]\n", + ")\n", + "\n", + "# Ask questions about AutoGen\n", + "stream = rag_assistant.run_stream(task=\"What is AgentChat?\")\n", + "await Console(stream)\n", + "\n", + "# Remember to close the memory when done\n", + "await rag_memory.close()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This implementation provides a RAG agent that can answer questions based on AutoGen documentation. When a question is asked, the Memory system retrieves relevant chunks and adds them to the context, enabling the assistant to generate informed responses.\n", + "\n", + "For production systems, you might want to:\n", + "1. Implement more sophisticated chunking strategies\n", + "2. Add metadata filtering capabilities\n", + "3. Customize the retrieval scoring\n", + "4. Optimize embedding models for your specific domain\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Mem0Memory Example\n", + "\n", + "`autogen_ext.memory.mem0.Mem0Memory` provides integration with `Mem0.ai`'s memory system. It supports both cloud-based and local backends, offering advanced memory capabilities for agents. The implementation handles proper retrieval and context updating, making it suitable for production environments.\n", + "\n", + "In the following example, we'll demonstrate how to use `Mem0Memory` to maintain persistent memories across conversations:" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_core.memory import MemoryContent, MemoryMimeType\n", + "from autogen_ext.memory.mem0 import Mem0Memory\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "# Initialize Mem0 cloud memory (requires API key)\n", + "# For local deployment, use is_cloud=False with appropriate config\n", + "mem0_memory = Mem0Memory(\n", + " is_cloud=True,\n", + " limit=5, # Maximum number of memories to retrieve\n", + ")\n", + "\n", + "# Add user preferences to memory\n", + "await mem0_memory.add(\n", + " MemoryContent(\n", + " content=\"The weather should be in metric units\",\n", + " mime_type=MemoryMimeType.TEXT,\n", + " metadata={\"category\": \"preferences\", \"type\": \"units\"},\n", + " )\n", + ")\n", + "\n", + "await mem0_memory.add(\n", + " MemoryContent(\n", + " content=\"Meal recipe must be vegan\",\n", + " mime_type=MemoryMimeType.TEXT,\n", + " metadata={\"category\": \"preferences\", \"type\": \"dietary\"},\n", + " )\n", + ")\n", + "\n", + "# Create assistant with mem0 memory\n", + "assistant_agent = AssistantAgent(\n", + " name=\"assistant_agent\",\n", + " model_client=OpenAIChatCompletionClient(\n", + " model=\"gpt-4o-2024-08-06\",\n", + " ),\n", + " tools=[get_weather],\n", + " memory=[mem0_memory],\n", + ")\n", + "\n", + "# Ask about the weather\n", + "stream = assistant_agent.run_stream(task=\"What are my dietary preferences?\")\n", + "await Console(stream)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The example above demonstrates how Mem0Memory can be used with an assistant agent. The memory integration ensures that:\n", + "\n", + "1. All agent interactions are stored in Mem0 for future reference\n", + "2. Relevant memories (like user preferences) are automatically retrieved and added to the context\n", + "3. The agent can maintain consistent behavior based on stored memories\n", + "\n", + "Mem0Memory is particularly useful for:\n", + "- Long-running agent deployments that need persistent memory\n", + "- Applications requiring enhanced privacy controls\n", + "- Teams wanting unified memory management across agents\n", + "- Use cases needing advanced memory filtering and analytics\n", + "\n", + "Just like ChromaDBVectorMemory, you can serialize Mem0Memory configurations:" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [], + "source": [ + "# Serialize the memory configuration\n", + "config_json = mem0_memory.dump_component().model_dump_json()\n", + "print(f\"Memory config JSON: {config_json[:100]}...\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "python", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.2" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/migration-guide.md b/python/docs/src/user-guide/agentchat-user-guide/migration-guide.md similarity index 97% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/migration-guide.md rename to python/docs/src/user-guide/agentchat-user-guide/migration-guide.md index d0533435fe42..ce93cee66ea5 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/migration-guide.md +++ b/python/docs/src/user-guide/agentchat-user-guide/migration-guide.md @@ -462,18 +462,18 @@ and implement the `on_messages`, `on_reset`, and `produced_message_types` method from typing import Sequence from autogen_core import CancellationToken from autogen_agentchat.agents import BaseChatAgent -from autogen_agentchat.messages import TextMessage, ChatMessage +from autogen_agentchat.messages import TextMessage, BaseChatMessage from autogen_agentchat.base import Response class CustomAgent(BaseChatAgent): - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: return Response(chat_message=TextMessage(content="Custom reply", source=self.name)) async def on_reset(self, cancellation_token: CancellationToken) -> None: pass @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: return (TextMessage,) ``` @@ -691,7 +691,7 @@ async def main() -> None: if user_input == "exit": break response = await assistant.on_messages([TextMessage(content=user_input, source="user")], CancellationToken()) - print("Assistant:", response.chat_message.content) + print("Assistant:", response.chat_message.to_text()) await model_client.close() asyncio.run(main()) @@ -742,8 +742,8 @@ You can use the following conversion functions to convert between a v0.4 message from typing import Any, Dict, List, Literal from autogen_agentchat.messages import ( - AgentEvent, - ChatMessage, + BaseAgentEvent, + BaseChatMessage, HandoffMessage, MultiModalMessage, StopMessage, @@ -757,14 +757,14 @@ from autogen_core.models import FunctionExecutionResult def convert_to_v02_message( - message: AgentEvent | ChatMessage, + message: BaseAgentEvent | BaseChatMessage, role: Literal["assistant", "user", "tool"], image_detail: Literal["auto", "high", "low"] = "auto", ) -> Dict[str, Any]: """Convert a v0.4 AgentChat message to a v0.2 message. Args: - message (AgentEvent | ChatMessage): The message to convert. + message (BaseAgentEvent | BaseChatMessage): The message to convert. role (Literal["assistant", "user", "tool"]): The role of the message. image_detail (Literal["auto", "high", "low"], optional): The detail level of image content in multi-modal message. Defaults to "auto". @@ -810,7 +810,7 @@ def convert_to_v02_message( return v02_message -def convert_to_v04_message(message: Dict[str, Any]) -> AgentEvent | ChatMessage: +def convert_to_v04_message(message: Dict[str, Any]) -> BaseAgentEvent | BaseChatMessage: """Convert a v0.2 message to a v0.4 AgentChat message.""" if "tool_calls" in message: tool_calls: List[FunctionCall] = [] @@ -1065,7 +1065,7 @@ import asyncio from typing import Sequence from autogen_agentchat.agents import AssistantAgent from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination -from autogen_agentchat.messages import AgentEvent, ChatMessage +from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage from autogen_agentchat.teams import SelectorGroupChat from autogen_agentchat.ui import Console from autogen_ext.models.openai import OpenAIChatCompletionClient @@ -1141,7 +1141,7 @@ def create_team(model_client : OpenAIChatCompletionClient) -> SelectorGroupChat: # The selector function is a function that takes the current message thread of the group chat # and returns the next speaker's name. If None is returned, the LLM-based selection method will be used. - def selector_func(messages: Sequence[AgentEvent | ChatMessage]) -> str | None: + def selector_func(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> str | None: if messages[-1].source != planning_agent.name: return planning_agent.name # Always return to the planning agent after the other agents have spoken. return None @@ -1190,12 +1190,12 @@ from typing import Sequence from autogen_core import CancellationToken from autogen_agentchat.agents import BaseChatAgent from autogen_agentchat.teams import RoundRobinGroupChat -from autogen_agentchat.messages import TextMessage, ChatMessage +from autogen_agentchat.messages import TextMessage, BaseChatMessage from autogen_agentchat.base import Response class CountingAgent(BaseChatAgent): """An agent that returns a new number by adding 1 to the last number in the input messages.""" - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: if len(messages) == 0: last_number = 0 # Start from 0 if no messages are given. else: @@ -1207,7 +1207,7 @@ class CountingAgent(BaseChatAgent): pass @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: return (TextMessage,) class NestedCountingAgent(BaseChatAgent): @@ -1217,7 +1217,7 @@ class NestedCountingAgent(BaseChatAgent): super().__init__(name, description="An agent that counts numbers.") self._counting_team = counting_team - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: # Run the inner team with the given messages and returns the last message produced by the team. result = await self._counting_team.run(task=messages, cancellation_token=cancellation_token) # To stream the inner messages, implement `on_messages_stream` and use that to implement `on_messages`. @@ -1229,7 +1229,7 @@ class NestedCountingAgent(BaseChatAgent): await self._counting_team.reset() @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: return (TextMessage,) async def main() -> None: @@ -1331,7 +1331,7 @@ async def main() -> None: if user_input == "exit": break response = await assistant.on_messages([TextMessage(content=user_input, source="user")], CancellationToken()) - print("Assistant:", response.chat_message.content) + print("Assistant:", response.chat_message.to_text()) await model_client.close() diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/quickstart.ipynb b/python/docs/src/user-guide/agentchat-user-guide/quickstart.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/quickstart.ipynb rename to python/docs/src/user-guide/agentchat-user-guide/quickstart.ipynb diff --git a/python/docs/src/user-guide/agentchat-user-guide/selector-group-chat.ipynb b/python/docs/src/user-guide/agentchat-user-guide/selector-group-chat.ipynb new file mode 100644 index 000000000000..64ce93a93aa4 --- /dev/null +++ b/python/docs/src/user-guide/agentchat-user-guide/selector-group-chat.ipynb @@ -0,0 +1,1035 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Selector Group Chat" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "{py:class}`~autogen_agentchat.teams.SelectorGroupChat` implements a team where participants take turns broadcasting messages to all other members. A generative model (e.g., an LLM) selects the next speaker based on the shared context, enabling dynamic, context-aware collaboration.\n", + "\n", + "Key features include:\n", + "\n", + "- Model-based speaker selection\n", + "- Configurable participant roles and descriptions\n", + "- Prevention of consecutive turns by the same speaker (optional)\n", + "- Customizable selection prompting\n", + "- Customizable selection function to override the default model-based selection\n", + "- Customizable candidate function to narrow-down the set of agents for selection using model\n", + "\n", + "```{note}\n", + "{py:class}`~autogen_agentchat.teams.SelectorGroupChat` is a high-level API. For more control and customization, refer to the [Group Chat Pattern](../core-user-guide/design-patterns/group-chat.ipynb) in the Core API documentation to implement your own group chat logic.\n", + "```\n", + "\n", + "## How Does it Work?\n", + "\n", + "{py:class}`~autogen_agentchat.teams.SelectorGroupChat` is a group chat similar to {py:class}`~autogen_agentchat.teams.RoundRobinGroupChat`,\n", + "but with a model-based next speaker selection mechanism.\n", + "When the team receives a task through {py:meth}`~autogen_agentchat.teams.BaseGroupChat.run` or {py:meth}`~autogen_agentchat.teams.BaseGroupChat.run_stream`,\n", + "the following steps are executed:\n", + "\n", + "1. The team analyzes the current conversation context, including the conversation history and participants' {py:attr}`~autogen_agentchat.base.ChatAgent.name` and {py:attr}`~autogen_agentchat.base.ChatAgent.description` attributes, to determine the next speaker using a model. By default, the team will not select the same speak consecutively unless it is the only agent available. This can be changed by setting `allow_repeated_speaker=True`. You can also override the model by providing a custom selection function.\n", + "2. The team prompts the selected speaker agent to provide a response, which is then **broadcasted** to all other participants.\n", + "3. The termination condition is checked to determine if the conversation should end, if not, the process repeats from step 1.\n", + "4. When the conversation ends, the team returns the {py:class}`~autogen_agentchat.base.TaskResult` containing the conversation history from this task.\n", + "\n", + "Once the team finishes the task, the conversation context is kept within the team and all participants, so the next task can continue from the previous conversation context.\n", + "You can reset the conversation context by calling {py:meth}`~autogen_agentchat.teams.BaseGroupChat.reset`.\n", + "\n", + "In this section, we will demonstrate how to use {py:class}`~autogen_agentchat.teams.SelectorGroupChat` with a simple example for a web search and data analysis task." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Example: Web Search/Analysis" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from typing import List, Sequence\n", + "\n", + "from autogen_agentchat.agents import AssistantAgent, UserProxyAgent\n", + "from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination\n", + "from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage\n", + "from autogen_agentchat.teams import SelectorGroupChat\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Agents\n", + "\n", + "![Selector Group Chat](selector-group-chat.svg)\n", + "\n", + "This system uses three specialized agents:\n", + "\n", + "- **Planning Agent**: The strategic coordinator that breaks down complex tasks into manageable subtasks. \n", + "- **Web Search Agent**: An information retrieval specialist that interfaces with the `search_web_tool`.\n", + "- **Data Analyst Agent**: An agent specialist in performing calculations equipped with `percentage_change_tool`. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The tools `search_web_tool` and `percentage_change_tool` are external tools that the agents can use to perform their tasks." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "# Note: This example uses mock tools instead of real APIs for demonstration purposes\n", + "def search_web_tool(query: str) -> str:\n", + " if \"2006-2007\" in query:\n", + " return \"\"\"Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", + " Udonis Haslem: 844 points\n", + " Dwayne Wade: 1397 points\n", + " James Posey: 550 points\n", + " ...\n", + " \"\"\"\n", + " elif \"2007-2008\" in query:\n", + " return \"The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\"\n", + " elif \"2008-2009\" in query:\n", + " return \"The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\"\n", + " return \"No data found.\"\n", + "\n", + "\n", + "def percentage_change_tool(start: float, end: float) -> float:\n", + " return ((end - start) / start) * 100" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Let's create the specialized agents using the {py:class}`~autogen_agentchat.agents.AssistantAgent` class.\n", + "It is important to note that the agents' {py:attr}`~autogen_agentchat.base.ChatAgent.name` and {py:attr}`~autogen_agentchat.base.ChatAgent.description` attributes are used by the model to determine the next speaker,\n", + "so it is recommended to provide meaningful names and descriptions." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "model_client = OpenAIChatCompletionClient(model=\"gpt-4o\")\n", + "\n", + "planning_agent = AssistantAgent(\n", + " \"PlanningAgent\",\n", + " description=\"An agent for planning tasks, this agent should be the first to engage when given a new task.\",\n", + " model_client=model_client,\n", + " system_message=\"\"\"\n", + " You are a planning agent.\n", + " Your job is to break down complex tasks into smaller, manageable subtasks.\n", + " Your team members are:\n", + " WebSearchAgent: Searches for information\n", + " DataAnalystAgent: Performs calculations\n", + "\n", + " You only plan and delegate tasks - you do not execute them yourself.\n", + "\n", + " When assigning tasks, use this format:\n", + " 1. : \n", + "\n", + " After all tasks are complete, summarize the findings and end with \"TERMINATE\".\n", + " \"\"\",\n", + ")\n", + "\n", + "web_search_agent = AssistantAgent(\n", + " \"WebSearchAgent\",\n", + " description=\"An agent for searching information on the web.\",\n", + " tools=[search_web_tool],\n", + " model_client=model_client,\n", + " system_message=\"\"\"\n", + " You are a web search agent.\n", + " Your only tool is search_tool - use it to find information.\n", + " You make only one search call at a time.\n", + " Once you have the results, you never do calculations based on them.\n", + " \"\"\",\n", + ")\n", + "\n", + "data_analyst_agent = AssistantAgent(\n", + " \"DataAnalystAgent\",\n", + " description=\"An agent for performing calculations.\",\n", + " model_client=model_client,\n", + " tools=[percentage_change_tool],\n", + " system_message=\"\"\"\n", + " You are a data analyst.\n", + " Given the tasks you have been assigned, you should analyze the data and provide results using the tools provided.\n", + " If you have not seen the data, ask for it.\n", + " \"\"\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "```{note}\n", + "By default, {py:class}`~autogen_agentchat.agents.AssistantAgent` returns the\n", + "tool output as the response. If your tool does not return a well-formed\n", + "string in natural language format, you may want to add a reflection step\n", + "within the agent by setting `reflect_on_tool_use=True` when creating the agent.\n", + "This will allow the agent to reflect on the tool output and provide a natural\n", + "language response.\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Workflow\n", + "\n", + "1. The task is received by the {py:class}`~autogen_agentchat.teams.SelectorGroupChat` which, based on agent descriptions, selects the most appropriate agent to handle the initial task (typically the Planning Agent).\n", + "\n", + "2. The **Planning Agent** analyzes the task and breaks it down into subtasks, assigning each to the most appropriate agent using the format:\n", + " ` : `\n", + "\n", + "3. Based on the conversation context and agent descriptions, the {py:class}`~autogen_agent.teams.SelectorGroupChat` manager dynamically selects the next agent to handle their assigned subtask.\n", + "\n", + "4. The **Web Search Agent** performs searches one at a time, storing results in the shared conversation history.\n", + "\n", + "5. The **Data Analyst** processes the gathered information using available calculation tools when selected.\n", + "\n", + "6. The workflow continues with agents being dynamically selected until either:\n", + " - The Planning Agent determines all subtasks are complete and sends \"TERMINATE\"\n", + " - An alternative termination condition is met (e.g., a maximum number of messages)\n", + "\n", + "When defining your agents, make sure to include a helpful {py:attr}`~autogen_agentchat.base.ChatAgent.description` since this is used to decide which agent to select next." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Termination Conditions\n", + "\n", + "Let's use two termination conditions:\n", + "{py:class}`~autogen_agentchat.conditions.TextMentionTermination` to end the conversation when the Planning Agent sends \"TERMINATE\",\n", + "and {py:class}`~autogen_agentchat.conditions.MaxMessageTermination` to limit the conversation to 25 messages to avoid infinite loop." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "text_mention_termination = TextMentionTermination(\"TERMINATE\")\n", + "max_messages_termination = MaxMessageTermination(max_messages=25)\n", + "termination = text_mention_termination | max_messages_termination" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Selector Prompt\n", + "\n", + "{py:class}`~autogen_agentchat.teams.SelectorGroupChat` uses a model to select\n", + "the next speaker based on the conversation context.\n", + "We will use a custom selector prompt to properly align with the workflow." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "selector_prompt = \"\"\"Select an agent to perform task.\n", + "\n", + "{roles}\n", + "\n", + "Current conversation context:\n", + "{history}\n", + "\n", + "Read the above conversation, then select an agent from {participants} to perform the next task.\n", + "Make sure the planner agent has assigned tasks before other agents start working.\n", + "Only select one agent.\n", + "\"\"\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The string variables available in the selector prompt are:\n", + "- `{participants}`: The names of candidates for selection. The format is `[\"\", \"\", ...]`.\n", + "- `{roles}`: A newline-separated list of names and descriptions of the candidate agents. The format for each line is: `\" : \"`.\n", + "- `{history}`: The conversation history formatted as a double newline separated of names and message content. The format for each message is: `\" : \"`.\n", + "\n", + "```{tip}\n", + "Try not to overload the model with too much instruction in the selector prompt.\n", + "\n", + "What is too much? It depends on the capabilities of the model you are using.\n", + "For GPT-4o and equivalents, you can use a selector prompt with a condition for when each speaker should be selected.\n", + "For smaller models such as Phi-4, you should keep the selector prompt as simple as possible\n", + "such as the one used in this example.\n", + "\n", + "Generally, if you find yourself writing multiple conditions for each agent,\n", + "it is a sign that you should consider using a custom selection function,\n", + "or breaking down the task into smaller, sequential tasks to be handled by\n", + "separate agents or teams.\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Running the Team\n", + "\n", + "Let's create the team with the agents, termination conditions, and custom selector prompt." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "team = SelectorGroupChat(\n", + " [planning_agent, web_search_agent, data_analyst_agent],\n", + " model_client=model_client,\n", + " termination_condition=termination,\n", + " selector_prompt=selector_prompt,\n", + " allow_repeated_speaker=True, # Allow an agent to speak multiple turns in a row.\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we run the team with a task to find information about an NBA player." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "task = \"Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\"" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n", + "---------- PlanningAgent ----------\n", + "To complete this task, we need to perform the following subtasks:\n", + "\n", + "1. Find out which Miami Heat player had the highest points in the 2006-2007 season.\n", + "2. Gather data on this player's total rebounds for the 2007-2008 season.\n", + "3. Gather data on this player's total rebounds for the 2008-2009 season.\n", + "4. Calculate the percentage change in the player's total rebounds between the 2007-2008 and 2008-2009 seasons.\n", + "\n", + "I'll assign these tasks accordingly:\n", + "\n", + "1. WebSearchAgent: Search for the Miami Heat player with the highest points in the 2006-2007 NBA season.\n", + "2. WebSearchAgent: Find the total rebounds for this player in the 2007-2008 NBA season.\n", + "3. WebSearchAgent: Find the total rebounds for this player in the 2008-2009 NBA season.\n", + "4. DataAnalystAgent: Calculate the percentage change in total rebounds from the 2007-2008 season to the 2008-2009 season for this player.\n", + "---------- WebSearchAgent ----------\n", + "[FunctionCall(id='call_89tUNHaAM0kKQYPJLleGUKK7', arguments='{\"query\":\"Miami Heat player highest points 2006-2007 season\"}', name='search_web_tool')]\n", + "---------- WebSearchAgent ----------\n", + "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', name='search_web_tool', call_id='call_89tUNHaAM0kKQYPJLleGUKK7', is_error=False)]\n", + "---------- WebSearchAgent ----------\n", + "Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", + " Udonis Haslem: 844 points\n", + " Dwayne Wade: 1397 points\n", + " James Posey: 550 points\n", + " ...\n", + " \n", + "---------- WebSearchAgent ----------\n", + "The Miami Heat player with the highest points in the 2006-2007 season was Dwyane Wade, with 1,397 points.\n", + "\n", + "Next, I will search for Dwyane Wade's total rebounds for the 2007-2008 season.\n", + "---------- WebSearchAgent ----------\n", + "[FunctionCall(id='call_RC55TkSjG3JXRuVOTPrcE1RL', arguments='{\"query\":\"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool')]\n", + "---------- WebSearchAgent ----------\n", + "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', name='search_web_tool', call_id='call_RC55TkSjG3JXRuVOTPrcE1RL', is_error=False)]\n", + "---------- WebSearchAgent ----------\n", + "The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n", + "---------- WebSearchAgent ----------\n", + "[FunctionCall(id='call_pBXoABrErDow0rZjw3tjOZol', arguments='{\"query\":\"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')]\n", + "---------- WebSearchAgent ----------\n", + "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', name='search_web_tool', call_id='call_pBXoABrErDow0rZjw3tjOZol', is_error=False)]\n", + "---------- WebSearchAgent ----------\n", + "The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n", + "---------- DataAnalystAgent ----------\n", + "[FunctionCall(id='call_qMxxXtcJsiK8KFSSCx3zm0is', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')]\n", + "---------- DataAnalystAgent ----------\n", + "[FunctionExecutionResult(content='85.98130841121495', name='percentage_change_tool', call_id='call_qMxxXtcJsiK8KFSSCx3zm0is', is_error=False)]\n", + "---------- DataAnalystAgent ----------\n", + "85.98130841121495\n", + "---------- PlanningAgent ----------\n", + "The player with the highest points for the Miami Heat in the 2006-2007 NBA season was Dwyane Wade, who scored 1,397 points. The percentage change in Dwyane Wade's total rebounds from 214 in the 2007-2008 season to 398 in the 2008-2009 season is approximately 85.98%.\n", + "\n", + "TERMINATE\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=161, completion_tokens=220), metadata={}, content=\"To complete this task, we need to perform the following subtasks:\\n\\n1. Find out which Miami Heat player had the highest points in the 2006-2007 season.\\n2. Gather data on this player's total rebounds for the 2007-2008 season.\\n3. Gather data on this player's total rebounds for the 2008-2009 season.\\n4. Calculate the percentage change in the player's total rebounds between the 2007-2008 and 2008-2009 seasons.\\n\\nI'll assign these tasks accordingly:\\n\\n1. WebSearchAgent: Search for the Miami Heat player with the highest points in the 2006-2007 NBA season.\\n2. WebSearchAgent: Find the total rebounds for this player in the 2007-2008 NBA season.\\n3. WebSearchAgent: Find the total rebounds for this player in the 2008-2009 NBA season.\\n4. DataAnalystAgent: Calculate the percentage change in total rebounds from the 2007-2008 season to the 2008-2009 season for this player.\", type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=368, completion_tokens=27), metadata={}, content=[FunctionCall(id='call_89tUNHaAM0kKQYPJLleGUKK7', arguments='{\"query\":\"Miami Heat player highest points 2006-2007 season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', name='search_web_tool', call_id='call_89tUNHaAM0kKQYPJLleGUKK7', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, metadata={}, content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', type='ToolCallSummaryMessage'), ThoughtEvent(source='WebSearchAgent', models_usage=None, metadata={}, content=\"The Miami Heat player with the highest points in the 2006-2007 season was Dwyane Wade, with 1,397 points.\\n\\nNext, I will search for Dwyane Wade's total rebounds for the 2007-2008 season.\", type='ThoughtEvent'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=460, completion_tokens=83), metadata={}, content=[FunctionCall(id='call_RC55TkSjG3JXRuVOTPrcE1RL', arguments='{\"query\":\"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', name='search_web_tool', call_id='call_RC55TkSjG3JXRuVOTPrcE1RL', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, metadata={}, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', type='ToolCallSummaryMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=585, completion_tokens=28), metadata={}, content=[FunctionCall(id='call_pBXoABrErDow0rZjw3tjOZol', arguments='{\"query\":\"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', name='search_web_tool', call_id='call_pBXoABrErDow0rZjw3tjOZol', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, metadata={}, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', type='ToolCallSummaryMessage'), ToolCallRequestEvent(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=496, completion_tokens=21), metadata={}, content=[FunctionCall(id='call_qMxxXtcJsiK8KFSSCx3zm0is', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='DataAnalystAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='85.98130841121495', name='percentage_change_tool', call_id='call_qMxxXtcJsiK8KFSSCx3zm0is', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='DataAnalystAgent', models_usage=None, metadata={}, content='85.98130841121495', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=528, completion_tokens=80), metadata={}, content=\"The player with the highest points for the Miami Heat in the 2006-2007 NBA season was Dwyane Wade, who scored 1,397 points. The percentage change in Dwyane Wade's total rebounds from 214 in the 2007-2008 season to 398 in the 2008-2009 season is approximately 85.98%.\\n\\nTERMINATE\", type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Use asyncio.run(...) if you are running this in a script.\n", + "await Console(team.run_stream(task=task))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "As we can see, after the Web Search Agent conducts the necessary searches and the Data Analyst Agent completes the necessary calculations, we find that Dwayne Wade was the Miami Heat player with the highest points in the 2006-2007 season, and the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons is 85.98%!" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Custom Selector Function" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Often times we want better control over the selection process.\n", + "To this end, we can set the `selector_func` argument with a custom selector function to override the default model-based selection.\n", + "This allows us to implement more complex selection logic and state-based transitions.\n", + "\n", + "For instance, we want the Planning Agent to speak immediately after any specialized agent to check the progress.\n", + "\n", + "```{note}\n", + "Returning `None` from the custom selector function will use the default model-based selection.\n", + "``` \n", + "\n", + "```{note}\n", + "Custom selector functions are not [serialized](https://microsoft.github.io/autogen/stable/user-guide/agentchat-user-guide/serialize-components.html) when `.dump_component()` is called on the SelectorGroupChat team . If you need to serialize team configurations with custom selector functions, consider implementing custom workflows and serialization logic.\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n", + "---------- PlanningAgent ----------\n", + "To answer this question, we need to follow these steps: \n", + "\n", + "1. Identify the Miami Heat player with the highest points in the 2006-2007 season.\n", + "2. Retrieve the total rebounds of that player for the 2007-2008 and 2008-2009 seasons.\n", + "3. Calculate the percentage change in his total rebounds between the two seasons.\n", + "\n", + "Let's delegate these tasks:\n", + "\n", + "1. WebSearchAgent: Find the Miami Heat player with the highest points in the 2006-2007 NBA season.\n", + "2. WebSearchAgent: Retrieve the total rebounds for the identified player during the 2007-2008 NBA season.\n", + "3. WebSearchAgent: Retrieve the total rebounds for the identified player during the 2008-2009 NBA season.\n", + "4. DataAnalystAgent: Calculate the percentage change in total rebounds between the 2007-2008 and 2008-2009 seasons for the player found.\n", + "---------- WebSearchAgent ----------\n", + "[FunctionCall(id='call_Pz82ndNLSV4cH0Sg6g7ArP4L', arguments='{\"query\":\"Miami Heat player highest points 2006-2007 season\"}', name='search_web_tool')]\n", + "---------- WebSearchAgent ----------\n", + "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', call_id='call_Pz82ndNLSV4cH0Sg6g7ArP4L')]\n", + "---------- WebSearchAgent ----------\n", + "Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", + " Udonis Haslem: 844 points\n", + " Dwayne Wade: 1397 points\n", + " James Posey: 550 points\n", + " ...\n", + " \n", + "---------- PlanningAgent ----------\n", + "Great! Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season. Now, let's continue with the next tasks:\n", + "\n", + "2. WebSearchAgent: Retrieve the total rebounds for Dwyane Wade during the 2007-2008 NBA season.\n", + "3. WebSearchAgent: Retrieve the total rebounds for Dwyane Wade during the 2008-2009 NBA season.\n", + "---------- WebSearchAgent ----------\n", + "[FunctionCall(id='call_3qv9so2DXFZIHtzqDIfXoFID', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool'), FunctionCall(id='call_Vh7zzzWUeiUAvaYjP0If0k1k', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')]\n", + "---------- WebSearchAgent ----------\n", + "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_3qv9so2DXFZIHtzqDIfXoFID'), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_Vh7zzzWUeiUAvaYjP0If0k1k')]\n", + "---------- WebSearchAgent ----------\n", + "The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n", + "The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n", + "---------- PlanningAgent ----------\n", + "Now let's calculate the percentage change in total rebounds between the 2007-2008 and 2008-2009 seasons for Dwyane Wade.\n", + "\n", + "4. DataAnalystAgent: Calculate the percentage change in total rebounds for Dwyane Wade between the 2007-2008 and 2008-2009 seasons.\n", + "---------- DataAnalystAgent ----------\n", + "[FunctionCall(id='call_FXnPSr6JVGfAWs3StIizbt2V', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')]\n", + "---------- DataAnalystAgent ----------\n", + "[FunctionExecutionResult(content='85.98130841121495', call_id='call_FXnPSr6JVGfAWs3StIizbt2V')]\n", + "---------- DataAnalystAgent ----------\n", + "85.98130841121495\n", + "---------- PlanningAgent ----------\n", + "Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring a total of 1397 points. The percentage change in his total rebounds from the 2007-2008 season (214 rebounds) to the 2008-2009 season (398 rebounds) is approximately 86.0%.\n", + "\n", + "TERMINATE\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=161, completion_tokens=192), content=\"To answer this question, we need to follow these steps: \\n\\n1. Identify the Miami Heat player with the highest points in the 2006-2007 season.\\n2. Retrieve the total rebounds of that player for the 2007-2008 and 2008-2009 seasons.\\n3. Calculate the percentage change in his total rebounds between the two seasons.\\n\\nLet's delegate these tasks:\\n\\n1. WebSearchAgent: Find the Miami Heat player with the highest points in the 2006-2007 NBA season.\\n2. WebSearchAgent: Retrieve the total rebounds for the identified player during the 2007-2008 NBA season.\\n3. WebSearchAgent: Retrieve the total rebounds for the identified player during the 2008-2009 NBA season.\\n4. DataAnalystAgent: Calculate the percentage change in total rebounds between the 2007-2008 and 2008-2009 seasons for the player found.\", type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=340, completion_tokens=27), content=[FunctionCall(id='call_Pz82ndNLSV4cH0Sg6g7ArP4L', arguments='{\"query\":\"Miami Heat player highest points 2006-2007 season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', call_id='call_Pz82ndNLSV4cH0Sg6g7ArP4L')], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=420, completion_tokens=87), content=\"Great! Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season. Now, let's continue with the next tasks:\\n\\n2. WebSearchAgent: Retrieve the total rebounds for Dwyane Wade during the 2007-2008 NBA season.\\n3. WebSearchAgent: Retrieve the total rebounds for Dwyane Wade during the 2008-2009 NBA season.\", type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=525, completion_tokens=71), content=[FunctionCall(id='call_3qv9so2DXFZIHtzqDIfXoFID', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool'), FunctionCall(id='call_Vh7zzzWUeiUAvaYjP0If0k1k', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_3qv9so2DXFZIHtzqDIfXoFID'), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_Vh7zzzWUeiUAvaYjP0If0k1k')], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\\nThe number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=569, completion_tokens=68), content=\"Now let's calculate the percentage change in total rebounds between the 2007-2008 and 2008-2009 seasons for Dwyane Wade.\\n\\n4. DataAnalystAgent: Calculate the percentage change in total rebounds for Dwyane Wade between the 2007-2008 and 2008-2009 seasons.\", type='TextMessage'), ToolCallRequestEvent(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=627, completion_tokens=21), content=[FunctionCall(id='call_FXnPSr6JVGfAWs3StIizbt2V', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='DataAnalystAgent', models_usage=None, content=[FunctionExecutionResult(content='85.98130841121495', call_id='call_FXnPSr6JVGfAWs3StIizbt2V')], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='DataAnalystAgent', models_usage=None, content='85.98130841121495', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=659, completion_tokens=76), content='Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring a total of 1397 points. The percentage change in his total rebounds from the 2007-2008 season (214 rebounds) to the 2008-2009 season (398 rebounds) is approximately 86.0%.\\n\\nTERMINATE', type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "def selector_func(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> str | None:\n", + " if messages[-1].source != planning_agent.name:\n", + " return planning_agent.name\n", + " return None\n", + "\n", + "\n", + "# Reset the previous team and run the chat again with the selector function.\n", + "await team.reset()\n", + "team = SelectorGroupChat(\n", + " [planning_agent, web_search_agent, data_analyst_agent],\n", + " model_client=model_client,\n", + " termination_condition=termination,\n", + " selector_prompt=selector_prompt,\n", + " allow_repeated_speaker=True,\n", + " selector_func=selector_func,\n", + ")\n", + "\n", + "await Console(team.run_stream(task=task))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can see from the conversation log that the Planning Agent always speaks immediately after the specialized agents.\n", + "\n", + "```{tip}\n", + "Each participant agent only makes one step (executing tools, generating a response, etc.)\n", + "on each turn. \n", + "If you want an {py:class}`~autogen_agentchat.agents.AssistantAgent` to repeat\n", + "until it stop returning a {py:class}`~autogen_agentchat.messages.ToolCallSummaryMessage`\n", + "when it has finished running all the tools it needs to run, you can do so by\n", + "checking the last message and returning the agent if it is a\n", + "{py:class}`~autogen_agentchat.messages.ToolCallSummaryMessage`.\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Custom Candidate Function" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "One more possible requirement might be to automatically select the next speaker from a filtered list of agents.\n", + "For this, we can set `candidate_func` parameter with a custom candidate function to filter down the list of potential agents for speaker selection for each turn of groupchat.\n", + "\n", + "This allow us to restrict speaker selection to a specific set of agents after a given agent.\n", + "\n", + "\n", + "```{note}\n", + "The `candidate_func` is only valid if `selector_func` is not set.\n", + "Returning `None` or an empty list `[]` from the custom candidate function will raise a `ValueError`.\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n", + "---------- PlanningAgent ----------\n", + "To answer this question, we'll break it down into two main subtasks:\n", + "\n", + "1. Identify the Miami Heat player with the highest points in the 2006-2007 season.\n", + "2. Calculate the percentage change in that player's total rebounds between the 2007-2008 and 2008-2009 seasons.\n", + "\n", + "Let's assign these tasks:\n", + "\n", + "1. WebSearchAgent: Search for the Miami Heat player with the highest points in the 2006-2007 NBA season.\n", + "2. WebSearchAgent: Find the total rebound statistics for that identified player for both the 2007-2008 and 2008-2009 NBA seasons.\n", + "3. DataAnalystAgent: Calculate the percentage change in the player's total rebounds between the 2007-2008 and 2008-2009 seasons once the data is retrieved.\n", + "---------- WebSearchAgent ----------\n", + "[FunctionCall(id='call_WtR5KTfEIxs3jIO25gjAw7dF', arguments='{\"query\":\"Miami Heat highest points scorer 2006-2007 NBA season\"}', name='search_web_tool')]\n", + "---------- WebSearchAgent ----------\n", + "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', name='search_web_tool', call_id='call_WtR5KTfEIxs3jIO25gjAw7dF', is_error=False)]\n", + "---------- WebSearchAgent ----------\n", + "Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", + " Udonis Haslem: 844 points\n", + " Dwayne Wade: 1397 points\n", + " James Posey: 550 points\n", + " ...\n", + " \n", + "---------- DataAnalystAgent ----------\n", + "[FunctionCall(id='call_9HA3DEacUl4WuG2G2PtRkXAO', arguments='{\"start\": 432, \"end\": 527}', name='percentage_change_tool')]\n", + "---------- DataAnalystAgent ----------\n", + "[FunctionExecutionResult(content='21.99074074074074', name='percentage_change_tool', call_id='call_9HA3DEacUl4WuG2G2PtRkXAO', is_error=False)]\n", + "---------- DataAnalystAgent ----------\n", + "21.99074074074074\n", + "---------- PlanningAgent ----------\n", + "It seems we've missed some context there, so let's assign the subtasks again for clarity:\n", + "\n", + "Based on the search results, Dwyane Wade had the highest points for the Miami Heat in the 2006-2007 season with 1397 points.\n", + "\n", + "Now, let's find the necessary rebound statistics:\n", + "\n", + "2. WebSearchAgent: Find Dwyane Wade's total rebound statistics for both the 2007-2008 and 2008-2009 NBA seasons.\n", + "3. DataAnalystAgent: Once the data is retrieved, calculate the percentage change in Dwyane Wade's total rebounds between the 2007-2008 and 2008-2009 seasons.\n", + "---------- WebSearchAgent ----------\n", + "[FunctionCall(id='call_3i1wTDSjkGg6Ev8YKYWkZK55', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 NBA season\"}', name='search_web_tool'), FunctionCall(id='call_NRAs6jHxXRi8zsvpW5WlHAaU', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 NBA season\"}', name='search_web_tool')]\n", + "---------- WebSearchAgent ----------\n", + "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', name='search_web_tool', call_id='call_3i1wTDSjkGg6Ev8YKYWkZK55', is_error=False), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', name='search_web_tool', call_id='call_NRAs6jHxXRi8zsvpW5WlHAaU', is_error=False)]\n", + "---------- WebSearchAgent ----------\n", + "The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n", + "The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n", + "---------- PlanningAgent ----------\n", + "The total rebounds for Dwyane Wade in the 2007-2008 season were 214, and in the 2008-2009 season, they were 398.\n", + "\n", + "Now, let's calculate the percentage change.\n", + "\n", + "3. DataAnalystAgent: Calculate the percentage change in Dwyane Wade's total rebounds from the 2007-2008 season to the 2008-2009 season.\n", + "---------- DataAnalystAgent ----------\n", + "[FunctionCall(id='call_XECA7ezz7VIKbf8IbZYSCSpI', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')]\n", + "---------- DataAnalystAgent ----------\n", + "[FunctionExecutionResult(content='85.98130841121495', name='percentage_change_tool', call_id='call_XECA7ezz7VIKbf8IbZYSCSpI', is_error=False)]\n", + "---------- DataAnalystAgent ----------\n", + "85.98130841121495\n", + "---------- PlanningAgent ----------\n", + "The Miami Heat player with the highest points in the 2006-2007 season was Dwyane Wade, with 1397 points. The percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons was approximately 85.98%.\n", + "\n", + "TERMINATE\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=161, completion_tokens=169), metadata={}, content=\"To answer this question, we'll break it down into two main subtasks:\\n\\n1. Identify the Miami Heat player with the highest points in the 2006-2007 season.\\n2. Calculate the percentage change in that player's total rebounds between the 2007-2008 and 2008-2009 seasons.\\n\\nLet's assign these tasks:\\n\\n1. WebSearchAgent: Search for the Miami Heat player with the highest points in the 2006-2007 NBA season.\\n2. WebSearchAgent: Find the total rebound statistics for that identified player for both the 2007-2008 and 2008-2009 NBA seasons.\\n3. DataAnalystAgent: Calculate the percentage change in the player's total rebounds between the 2007-2008 and 2008-2009 seasons once the data is retrieved.\", type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=324, completion_tokens=28), metadata={}, content=[FunctionCall(id='call_WtR5KTfEIxs3jIO25gjAw7dF', arguments='{\"query\":\"Miami Heat highest points scorer 2006-2007 NBA season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', name='search_web_tool', call_id='call_WtR5KTfEIxs3jIO25gjAw7dF', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, metadata={}, content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', type='ToolCallSummaryMessage'), ToolCallRequestEvent(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=390, completion_tokens=37), metadata={}, content=[FunctionCall(id='call_9HA3DEacUl4WuG2G2PtRkXAO', arguments='{\"start\": 432, \"end\": 527}', name='percentage_change_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='DataAnalystAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='21.99074074074074', name='percentage_change_tool', call_id='call_9HA3DEacUl4WuG2G2PtRkXAO', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='DataAnalystAgent', models_usage=None, metadata={}, content='21.99074074074074', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=413, completion_tokens=137), metadata={}, content=\"It seems we've missed some context there, so let's assign the subtasks again for clarity:\\n\\nBased on the search results, Dwyane Wade had the highest points for the Miami Heat in the 2006-2007 season with 1397 points.\\n\\nNow, let's find the necessary rebound statistics:\\n\\n2. WebSearchAgent: Find Dwyane Wade's total rebound statistics for both the 2007-2008 and 2008-2009 NBA seasons.\\n3. DataAnalystAgent: Once the data is retrieved, calculate the percentage change in Dwyane Wade's total rebounds between the 2007-2008 and 2008-2009 seasons.\", type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=576, completion_tokens=73), metadata={}, content=[FunctionCall(id='call_3i1wTDSjkGg6Ev8YKYWkZK55', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 NBA season\"}', name='search_web_tool'), FunctionCall(id='call_NRAs6jHxXRi8zsvpW5WlHAaU', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 NBA season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', name='search_web_tool', call_id='call_3i1wTDSjkGg6Ev8YKYWkZK55', is_error=False), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', name='search_web_tool', call_id='call_NRAs6jHxXRi8zsvpW5WlHAaU', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, metadata={}, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\\nThe number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=612, completion_tokens=84), metadata={}, content=\"The total rebounds for Dwyane Wade in the 2007-2008 season were 214, and in the 2008-2009 season, they were 398.\\n\\nNow, let's calculate the percentage change.\\n\\n3. DataAnalystAgent: Calculate the percentage change in Dwyane Wade's total rebounds from the 2007-2008 season to the 2008-2009 season.\", type='TextMessage'), ToolCallRequestEvent(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=720, completion_tokens=21), metadata={}, content=[FunctionCall(id='call_XECA7ezz7VIKbf8IbZYSCSpI', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='DataAnalystAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='85.98130841121495', name='percentage_change_tool', call_id='call_XECA7ezz7VIKbf8IbZYSCSpI', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='DataAnalystAgent', models_usage=None, metadata={}, content='85.98130841121495', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=718, completion_tokens=63), metadata={}, content='The Miami Heat player with the highest points in the 2006-2007 season was Dwyane Wade, with 1397 points. The percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons was approximately 85.98%.\\n\\nTERMINATE', type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "def candidate_func(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str]:\n", + " # keep planning_agent first one to plan out the tasks\n", + " if messages[-1].source == \"user\":\n", + " return [planning_agent.name]\n", + "\n", + " # if previous agent is planning_agent and if it explicitely asks for web_search_agent\n", + " # or data_analyst_agent or both (in-case of re-planning or re-assignment of tasks)\n", + " # then return those specific agents\n", + " last_message = messages[-1]\n", + " if last_message.source == planning_agent.name:\n", + " participants = []\n", + " if web_search_agent.name in last_message.to_text():\n", + " participants.append(web_search_agent.name)\n", + " if data_analyst_agent.name in last_message.to_text():\n", + " participants.append(data_analyst_agent.name)\n", + " if participants:\n", + " return participants # SelectorGroupChat will select from the remaining two agents.\n", + "\n", + " # we can assume that the task is finished once the web_search_agent\n", + " # and data_analyst_agent have took their turns, thus we send\n", + " # in planning_agent to terminate the chat\n", + " previous_set_of_agents = set(message.source for message in messages)\n", + " if web_search_agent.name in previous_set_of_agents and data_analyst_agent.name in previous_set_of_agents:\n", + " return [planning_agent.name]\n", + "\n", + " # if no-conditions are met then return all the agents\n", + " return [planning_agent.name, web_search_agent.name, data_analyst_agent.name]\n", + "\n", + "\n", + "# Reset the previous team and run the chat again with the selector function.\n", + "await team.reset()\n", + "team = SelectorGroupChat(\n", + " [planning_agent, web_search_agent, data_analyst_agent],\n", + " model_client=model_client,\n", + " termination_condition=termination,\n", + " candidate_func=candidate_func,\n", + ")\n", + "\n", + "await Console(team.run_stream(task=task))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can see from the conversation log that the Planning Agent returns to conversation once the Web Search Agent and Data Analyst Agent took their turns and it finds that the task was not finished as expected so it called the WebSearchAgent again to get rebound values and then called DataAnalysetAgent to get the percentage change." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## User Feedback\n", + "\n", + "We can add {py:class}`~autogen_agentchat.agents.UserProxyAgent` to the team to\n", + "provide user feedback during a run.\n", + "See [Human-in-the-Loop](./tutorial/human-in-the-loop.ipynb) for more details\n", + "about {py:class}`~autogen_agentchat.agents.UserProxyAgent`.\n", + "\n", + "To use the {py:class}`~autogen_agentchat.agents.UserProxyAgent` in the \n", + "web search example, we simply add it to the team and update the selector function\n", + "to always check for user feedback after the planning agent speaks.\n", + "If the user responds with `\"APPROVE\"`, the conversation continues, otherwise,\n", + "the planning agent tries again, until the user approves." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- PlanningAgent ----------\n", + "To address the user's query, we will need to perform the following tasks:\n", + "\n", + "1. Identify the Miami Heat player with the highest points in the 2006-2007 season.\n", + "2. Find the total rebounds for that player in the 2007-2008 season.\n", + "3. Find the total rebounds for that player in the 2008-2009 season.\n", + "4. Calculate the percentage change in the total rebounds between the 2007-2008 and 2008-2009 seasons.\n", + "\n", + "Let's assign these tasks:\n", + "\n", + "1. **WebSearchAgent**: Identify the Miami Heat player with the highest points in the 2006-2007 season.\n", + " \n", + "(Task 2 and 3 depend on the result of Task 1. We'll proceed with Tasks 2 and 3 once Task 1 is complete.)\n", + "---------- UserProxyAgent ----------\n", + "approve\n", + "---------- WebSearchAgent ----------\n", + "[FunctionCall(id='call_0prr3fUnG5CtisUG7QeygW0w', arguments='{\"query\":\"Miami Heat highest points scorer 2006-2007 NBA season\"}', name='search_web_tool')]\n", + "---------- WebSearchAgent ----------\n", + "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', call_id='call_0prr3fUnG5CtisUG7QeygW0w')]\n", + "---------- WebSearchAgent ----------\n", + "Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", + " Udonis Haslem: 844 points\n", + " Dwayne Wade: 1397 points\n", + " James Posey: 550 points\n", + " ...\n", + " \n", + "---------- PlanningAgent ----------\n", + "Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring 1397 points.\n", + "\n", + "Next, we need to find Dwyane Wade's total rebounds for the 2007-2008 and 2008-2009 seasons:\n", + "\n", + "2. **WebSearchAgent**: Find Dwyane Wade's total rebounds for the 2007-2008 season.\n", + "3. **WebSearchAgent**: Find Dwyane Wade's total rebounds for the 2008-2009 season.\n", + "---------- UserProxyAgent ----------\n", + "approve\n", + "---------- WebSearchAgent ----------\n", + "[FunctionCall(id='call_fBZe80NaBfruOVGwRWbhXyRm', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 NBA season\"}', name='search_web_tool'), FunctionCall(id='call_cURYibna4fGxySiL7IYt0c3s', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 NBA season\"}', name='search_web_tool')]\n", + "---------- WebSearchAgent ----------\n", + "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_fBZe80NaBfruOVGwRWbhXyRm'), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_cURYibna4fGxySiL7IYt0c3s')]\n", + "---------- WebSearchAgent ----------\n", + "The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n", + "The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n", + "---------- PlanningAgent ----------\n", + "Now that we have Dwyane Wade's total rebounds for both seasons, we can calculate the percentage change:\n", + "\n", + "4. **DataAnalystAgent**: Calculate the percentage change in Dwyane Wade's total rebounds from the 2007-2008 season (214 rebounds) to the 2008-2009 season (398 rebounds).\n", + "---------- UserProxyAgent ----------\n", + "approve\n", + "---------- DataAnalystAgent ----------\n", + "[FunctionCall(id='call_z3uog7t2x0z1Suzl5hACF9hY', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')]\n", + "---------- DataAnalystAgent ----------\n", + "[FunctionExecutionResult(content='85.98130841121495', call_id='call_z3uog7t2x0z1Suzl5hACF9hY')]\n", + "---------- DataAnalystAgent ----------\n", + "85.98130841121495\n", + "---------- PlanningAgent ----------\n", + "Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring 1397 points. His total rebounds increased from 214 in the 2007-2008 season to 398 in the 2008-2009 season, which is a percentage change of approximately 85.98%.\n", + "\n", + "TERMINATE\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=161, completion_tokens=166), content=\"To address the user's query, we will need to perform the following tasks:\\n\\n1. Identify the Miami Heat player with the highest points in the 2006-2007 season.\\n2. Find the total rebounds for that player in the 2007-2008 season.\\n3. Find the total rebounds for that player in the 2008-2009 season.\\n4. Calculate the percentage change in the total rebounds between the 2007-2008 and 2008-2009 seasons.\\n\\nLet's assign these tasks:\\n\\n1. **WebSearchAgent**: Identify the Miami Heat player with the highest points in the 2006-2007 season.\\n \\n(Task 2 and 3 depend on the result of Task 1. We'll proceed with Tasks 2 and 3 once Task 1 is complete.)\", type='TextMessage'), UserInputRequestedEvent(source='UserProxyAgent', models_usage=None, request_id='2a433f88-f886-4b39-a078-ea1acdcb2f9d', content='', type='UserInputRequestedEvent'), TextMessage(source='UserProxyAgent', models_usage=None, content='approve', type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=323, completion_tokens=28), content=[FunctionCall(id='call_0prr3fUnG5CtisUG7QeygW0w', arguments='{\"query\":\"Miami Heat highest points scorer 2006-2007 NBA season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', call_id='call_0prr3fUnG5CtisUG7QeygW0w')], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=403, completion_tokens=112), content=\"Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring 1397 points.\\n\\nNext, we need to find Dwyane Wade's total rebounds for the 2007-2008 and 2008-2009 seasons:\\n\\n2. **WebSearchAgent**: Find Dwyane Wade's total rebounds for the 2007-2008 season.\\n3. **WebSearchAgent**: Find Dwyane Wade's total rebounds for the 2008-2009 season.\", type='TextMessage'), UserInputRequestedEvent(source='UserProxyAgent', models_usage=None, request_id='23dd4570-2391-41e9-aeea-86598499792c', content='', type='UserInputRequestedEvent'), TextMessage(source='UserProxyAgent', models_usage=None, content='approve', type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=543, completion_tokens=73), content=[FunctionCall(id='call_fBZe80NaBfruOVGwRWbhXyRm', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 NBA season\"}', name='search_web_tool'), FunctionCall(id='call_cURYibna4fGxySiL7IYt0c3s', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 NBA season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_fBZe80NaBfruOVGwRWbhXyRm'), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_cURYibna4fGxySiL7IYt0c3s')], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\\nThe number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=586, completion_tokens=70), content=\"Now that we have Dwyane Wade's total rebounds for both seasons, we can calculate the percentage change:\\n\\n4. **DataAnalystAgent**: Calculate the percentage change in Dwyane Wade's total rebounds from the 2007-2008 season (214 rebounds) to the 2008-2009 season (398 rebounds).\", type='TextMessage'), UserInputRequestedEvent(source='UserProxyAgent', models_usage=None, request_id='e849d193-4ab3-4558-8560-7dbc062a0aee', content='', type='UserInputRequestedEvent'), TextMessage(source='UserProxyAgent', models_usage=None, content='approve', type='TextMessage'), ToolCallRequestEvent(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=655, completion_tokens=21), content=[FunctionCall(id='call_z3uog7t2x0z1Suzl5hACF9hY', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='DataAnalystAgent', models_usage=None, content=[FunctionExecutionResult(content='85.98130841121495', call_id='call_z3uog7t2x0z1Suzl5hACF9hY')], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='DataAnalystAgent', models_usage=None, content='85.98130841121495', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=687, completion_tokens=74), content='Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring 1397 points. His total rebounds increased from 214 in the 2007-2008 season to 398 in the 2008-2009 season, which is a percentage change of approximately 85.98%.\\n\\nTERMINATE', type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "user_proxy_agent = UserProxyAgent(\"UserProxyAgent\", description=\"A proxy for the user to approve or disapprove tasks.\")\n", + "\n", + "\n", + "def selector_func_with_user_proxy(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> str | None:\n", + " if messages[-1].source != planning_agent.name and messages[-1].source != user_proxy_agent.name:\n", + " # Planning agent should be the first to engage when given a new task, or check progress.\n", + " return planning_agent.name\n", + " if messages[-1].source == planning_agent.name:\n", + " if messages[-2].source == user_proxy_agent.name and \"APPROVE\" in messages[-1].content.upper(): # type: ignore\n", + " # User has approved the plan, proceed to the next agent.\n", + " return None\n", + " # Use the user proxy agent to get the user's approval to proceed.\n", + " return user_proxy_agent.name\n", + " if messages[-1].source == user_proxy_agent.name:\n", + " # If the user does not approve, return to the planning agent.\n", + " if \"APPROVE\" not in messages[-1].content.upper(): # type: ignore\n", + " return planning_agent.name\n", + " return None\n", + "\n", + "\n", + "# Reset the previous agents and run the chat again with the user proxy agent and selector function.\n", + "await team.reset()\n", + "team = SelectorGroupChat(\n", + " [planning_agent, web_search_agent, data_analyst_agent, user_proxy_agent],\n", + " model_client=model_client,\n", + " termination_condition=termination,\n", + " selector_prompt=selector_prompt,\n", + " selector_func=selector_func_with_user_proxy,\n", + " allow_repeated_speaker=True,\n", + ")\n", + "\n", + "await Console(team.run_stream(task=task))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now, the user's feedback is incorporated into the conversation flow,\n", + "and the user can approve or reject the planning agent's decisions." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Reasoning Models\n", + "\n", + "So far in the examples, we have used a `gpt-4o` model. Models like `gpt-4o`\n", + "and `gemini-1.5-flash` are great at following instructions, so you can\n", + "have relatively detailed instructions in the selector prompt for the team and the \n", + "system messages for each agent to guide their behavior.\n", + "\n", + "However, if you are using a reasoning model like `o3-mini`, you will need to\n", + "keep the selector prompt and system messages as simple and to the point as possible.\n", + "This is because the reasoning models are already good at coming up with their own \n", + "instructions given the context provided to them.\n", + "\n", + "This also means that we don't need a planning agent to break down the task\n", + "anymore, since the {py:class}`~autogen_agentchat.teams.SelectorGroupChat` that\n", + "uses a reasoning model can do that on its own.\n", + "\n", + "In the following example, we will use `o3-mini` as the model for the\n", + "agents and the team, and we will not use a planning agent.\n", + "Also, we are keeping the selector prompt and system messages as simple as possible." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [], + "source": [ + "model_client = OpenAIChatCompletionClient(model=\"o3-mini\")\n", + "\n", + "web_search_agent = AssistantAgent(\n", + " \"WebSearchAgent\",\n", + " description=\"An agent for searching information on the web.\",\n", + " tools=[search_web_tool],\n", + " model_client=model_client,\n", + " system_message=\"\"\"Use web search tool to find information.\"\"\",\n", + ")\n", + "\n", + "data_analyst_agent = AssistantAgent(\n", + " \"DataAnalystAgent\",\n", + " description=\"An agent for performing calculations.\",\n", + " model_client=model_client,\n", + " tools=[percentage_change_tool],\n", + " system_message=\"\"\"Use tool to perform calculation. If you have not seen the data, ask for it.\"\"\",\n", + ")\n", + "\n", + "user_proxy_agent = UserProxyAgent(\n", + " \"UserProxyAgent\",\n", + " description=\"A user to approve or disapprove tasks.\",\n", + ")\n", + "\n", + "selector_prompt = \"\"\"Select an agent to perform task.\n", + "\n", + "{roles}\n", + "\n", + "Current conversation context:\n", + "{history}\n", + "\n", + "Read the above conversation, then select an agent from {participants} to perform the next task.\n", + "When the task is complete, let the user approve or disapprove the task.\n", + "\"\"\"\n", + "\n", + "team = SelectorGroupChat(\n", + " [web_search_agent, data_analyst_agent, user_proxy_agent],\n", + " model_client=model_client,\n", + " termination_condition=termination, # Use the same termination condition as before.\n", + " selector_prompt=selector_prompt,\n", + " allow_repeated_speaker=True,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n", + "---------- WebSearchAgent ----------\n", + "[FunctionCall(id='call_hl7EP6Lp5jj5wEdxeNHTwUVG', arguments='{\"query\": \"Who was the Miami Heat player with the highest points in the 2006-2007 season Miami Heat statistics Dwyane Wade rebounds percentage change 2007-2008 2008-2009 seasons\"}', name='search_web_tool')]\n", + "---------- WebSearchAgent ----------\n", + "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', call_id='call_hl7EP6Lp5jj5wEdxeNHTwUVG', is_error=False)]\n", + "---------- WebSearchAgent ----------\n", + "Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", + " Udonis Haslem: 844 points\n", + " Dwayne Wade: 1397 points\n", + " James Posey: 550 points\n", + " ...\n", + " \n", + "---------- DataAnalystAgent ----------\n", + "I found that in the 2006–2007 season the player with the highest points was Dwyane Wade (with 1,397 points). Could you please provide Dwyane Wade’s total rebounds for the 2007–2008 and the 2008–2009 seasons so I can calculate the percentage change?\n", + "---------- WebSearchAgent ----------\n", + "[FunctionCall(id='call_lppGTILXDvO9waPwKO66ehK6', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 and 2008-2009 seasons for Miami Heat\"}', name='search_web_tool')]\n", + "---------- WebSearchAgent ----------\n", + "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_lppGTILXDvO9waPwKO66ehK6', is_error=False)]\n", + "---------- WebSearchAgent ----------\n", + "The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n", + "---------- DataAnalystAgent ----------\n", + "Could you please provide Dwyane Wade’s total rebounds in the 2008-2009 season?\n", + "---------- WebSearchAgent ----------\n", + "[FunctionCall(id='call_r8DBcbJtQfdtugLtyTrqOvoK', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season Miami Heat\"}', name='search_web_tool')]\n", + "---------- WebSearchAgent ----------\n", + "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_r8DBcbJtQfdtugLtyTrqOvoK', is_error=False)]\n", + "---------- WebSearchAgent ----------\n", + "The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n", + "---------- DataAnalystAgent ----------\n", + "[FunctionCall(id='call_4jejv1wM7V1osbBCxJze8aQM', arguments='{\"start\": 214, \"end\": 398}', name='percentage_change_tool')]\n", + "---------- DataAnalystAgent ----------\n", + "[FunctionExecutionResult(content='85.98130841121495', call_id='call_4jejv1wM7V1osbBCxJze8aQM', is_error=False)]\n", + "---------- DataAnalystAgent ----------\n", + "85.98130841121495\n", + "---------- DataAnalystAgent ----------\n", + "Dwyane Wade was the Miami Heat player with the highest total points (1,397) during the 2006-2007 season. His total rebounds increased by approximately 86% from 214 in the 2007-2008 season to 398 in the 2008-2009 season.\n", + "---------- UserProxyAgent ----------\n", + "Approve. TERMINATE\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?', type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=103, completion_tokens=384), content=[FunctionCall(id='call_hl7EP6Lp5jj5wEdxeNHTwUVG', arguments='{\"query\": \"Who was the Miami Heat player with the highest points in the 2006-2007 season Miami Heat statistics Dwyane Wade rebounds percentage change 2007-2008 2008-2009 seasons\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', call_id='call_hl7EP6Lp5jj5wEdxeNHTwUVG', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', type='ToolCallSummaryMessage'), TextMessage(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=183, completion_tokens=1038), content='I found that in the 2006–2007 season the player with the highest points was Dwyane Wade (with 1,397 points). Could you please provide Dwyane Wade’s total rebounds for the 2007–2008 and the 2008–2009 seasons so I can calculate the percentage change?', type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=299, completion_tokens=109), content=[FunctionCall(id='call_lppGTILXDvO9waPwKO66ehK6', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 and 2008-2009 seasons for Miami Heat\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_lppGTILXDvO9waPwKO66ehK6', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', type='ToolCallSummaryMessage'), TextMessage(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=291, completion_tokens=224), content='Could you please provide Dwyane Wade’s total rebounds in the 2008-2009 season?', type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=401, completion_tokens=37), content=[FunctionCall(id='call_r8DBcbJtQfdtugLtyTrqOvoK', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season Miami Heat\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_r8DBcbJtQfdtugLtyTrqOvoK', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', type='ToolCallSummaryMessage'), ToolCallRequestEvent(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=353, completion_tokens=158), content=[FunctionCall(id='call_4jejv1wM7V1osbBCxJze8aQM', arguments='{\"start\": 214, \"end\": 398}', name='percentage_change_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='DataAnalystAgent', models_usage=None, content=[FunctionExecutionResult(content='85.98130841121495', call_id='call_4jejv1wM7V1osbBCxJze8aQM', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='DataAnalystAgent', models_usage=None, content='85.98130841121495', type='ToolCallSummaryMessage'), TextMessage(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=394, completion_tokens=138), content='Dwyane Wade was the Miami Heat player with the highest total points (1,397) during the 2006-2007 season. His total rebounds increased by approximately 86% from 214 in the 2007-2008 season to 398 in the 2008-2009 season.', type='TextMessage'), UserInputRequestedEvent(source='UserProxyAgent', models_usage=None, request_id='b3b05408-73fc-47d4-b832-16c9f447cd6e', content='', type='UserInputRequestedEvent'), TextMessage(source='UserProxyAgent', models_usage=None, content='Approve. TERMINATE', type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "await Console(team.run_stream(task=task))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "```{tip}\n", + "For more guidance on how to prompt reasoning models, see the\n", + "Azure AI Services Blog on [Prompt Engineering for OpenAI's O1 and O3-mini Reasoning Models](https://techcommunity.microsoft.com/blog/azure-ai-services-blog/prompt-engineering-for-openai%E2%80%99s-o1-and-o3-mini-reasoning-models/4374010)\n", + "```" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/selector-group-chat.svg b/python/docs/src/user-guide/agentchat-user-guide/selector-group-chat.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/selector-group-chat.svg rename to python/docs/src/user-guide/agentchat-user-guide/selector-group-chat.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb b/python/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb similarity index 98% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb rename to python/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb index a6e41d6ef786..138ca8fb4601 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb +++ b/python/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb @@ -19,6 +19,10 @@ " \n", "```\n", "\n", + "```{note}\n", + "`selector_func` is not serializable and will be ignored during serialization and deserialization process.\n", + "```\n", + "\n", " \n", "### Termination Condition Example \n", "\n", diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/swarm.ipynb b/python/docs/src/user-guide/agentchat-user-guide/swarm.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/swarm.ipynb rename to python/docs/src/user-guide/agentchat-user-guide/swarm.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/swarm_customer_support.svg b/python/docs/src/user-guide/agentchat-user-guide/swarm_customer_support.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/swarm_customer_support.svg rename to python/docs/src/user-guide/agentchat-user-guide/swarm_customer_support.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/swarm_stock_research.svg b/python/docs/src/user-guide/agentchat-user-guide/swarm_stock_research.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/swarm_stock_research.svg rename to python/docs/src/user-guide/agentchat-user-guide/swarm_stock_research.svg diff --git a/python/docs/src/user-guide/agentchat-user-guide/tracing.ipynb b/python/docs/src/user-guide/agentchat-user-guide/tracing.ipynb new file mode 100644 index 000000000000..914fcafd4bb1 --- /dev/null +++ b/python/docs/src/user-guide/agentchat-user-guide/tracing.ipynb @@ -0,0 +1,331 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Tracing and Observability\n", + "\n", + "AutoGen has [built-in support for tracing](https://microsoft.github.io/autogen/dev/user-guide/core-user-guide/framework/telemetry.html) and observability for collecting comprehensive records on the execution of your application. This feature is useful for debugging, performance analysis, and understanding the flow of your application.\n", + "\n", + "This capability is powered by the [OpenTelemetry](https://opentelemetry.io/) library, which means you can use any OpenTelemetry-compatible backend to collect and analyze traces.\n", + "\n", + "AutoGen follows the [OpenTelemetry Semantic Conventions](https://opentelemetry.io/docs/specs/semconv/) for tracing, for agents and tools.\n", + "It also follows the [Semantic Conventions for GenAI Systems](https://opentelemetry.io/docs/specs/semconv/gen-ai/) currently under development.\n", + "\n", + "## Setup\n", + "\n", + "To begin, you need to install the OpenTelemetry Python package. You can do this using pip:\n", + "\n", + "```bash\n", + "pip install opentelemetry-sdk opentelemetry-exporter-otlp-proto-grpc opentelemetry-instrumentation-openai\n", + "```\n", + "\n", + "Once you have the SDK installed, the simplest way to set up tracing in AutoGen is to:\n", + "\n", + "1. Configure an OpenTelemetry tracer provider\n", + "2. Set up an exporter to send traces to your backend\n", + "3. Connect the tracer provider to the AutoGen runtime\n", + "\n", + "## Telemetry Backend\n", + "\n", + "To collect and view traces, you need to set up a telemetry backend. Several open-source options are available, including Jaeger, Zipkin. For this example, we will use Jaeger as our telemetry backend.\n", + "\n", + "For a quick start, you can run Jaeger locally using Docker:\n", + "\n", + "```bash\n", + "docker run -d --name jaeger \\\n", + " -e COLLECTOR_OTLP_ENABLED=true \\\n", + " -p 16686:16686 \\\n", + " -p 4317:4317 \\\n", + " -p 4318:4318 \\\n", + " jaegertracing/all-in-one:latest\n", + "```\n", + "\n", + "This command starts a Jaeger instance that listens on port 16686 for the Jaeger UI and port 4317 for the OpenTelemetry collector. You can access the Jaeger UI at `http://localhost:16686`.\n", + "\n", + "## Tracing an AgentChat Team\n", + "\n", + "In the following section, we will review how to enable tracing with an AutoGen GroupChat team. The AutoGen runtime already supports open telemetry (automatically logging message metadata). To begin, we will create a tracing service that will be used to instrument the AutoGen runtime. " + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Overriding of current TracerProvider is not allowed\n", + "Attempting to instrument while already instrumented\n" + ] + } + ], + "source": [ + "from opentelemetry import trace\n", + "from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter\n", + "from opentelemetry.instrumentation.openai import OpenAIInstrumentor\n", + "from opentelemetry.sdk.resources import Resource\n", + "from opentelemetry.sdk.trace import TracerProvider\n", + "from opentelemetry.sdk.trace.export import BatchSpanProcessor\n", + "\n", + "# Set up telemetry span exporter.\n", + "otel_exporter = OTLPSpanExporter(endpoint=\"http://localhost:4317\", insecure=True)\n", + "span_processor = BatchSpanProcessor(otel_exporter)\n", + "\n", + "# Set up telemetry trace provider.\n", + "tracer_provider = TracerProvider(resource=Resource({\"service.name\": \"autogen-test-agentchat\"}))\n", + "tracer_provider.add_span_processor(span_processor)\n", + "trace.set_tracer_provider(tracer_provider)\n", + "\n", + "# Instrument the OpenAI Python library\n", + "OpenAIInstrumentor().instrument()\n", + "\n", + "# we will get reference this tracer later using its service name\n", + "# tracer = trace.get_tracer(\"autogen-test-agentchat\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "All of the code to create a [team](./tutorial/teams.ipynb) should already be familiar to you.\n", + "\n", + "```{note}\n", + "AgentChat teams are run using the AutoGen Core's agent runtime.\n", + "In turn, the runtime is already instrumented to log, see [Core Telemetry Guide](../core-user-guide/framework/telemetry.md).\n", + "To disable the agent runtime telemetry, you can set the `trace_provider` to\n", + "`opentelemetry.trace.NoOpTracerProvider` in the runtime constructor.\n", + "\n", + "Additionally, you can set the environment variable `AUTOGEN_DISABLE_RUNTIME_TRACING` to `true` to disable the agent runtime telemetry if you don't have access to the runtime constructor. For example, if you are using `ComponentConfig`.\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination\n", + "from autogen_agentchat.teams import SelectorGroupChat\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_core import SingleThreadedAgentRuntime\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "\n", + "def search_web_tool(query: str) -> str:\n", + " if \"2006-2007\" in query:\n", + " return \"\"\"Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", + " Udonis Haslem: 844 points\n", + " Dwayne Wade: 1397 points\n", + " James Posey: 550 points\n", + " ...\n", + " \"\"\"\n", + " elif \"2007-2008\" in query:\n", + " return \"The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\"\n", + " elif \"2008-2009\" in query:\n", + " return \"The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\"\n", + " return \"No data found.\"\n", + "\n", + "\n", + "def percentage_change_tool(start: float, end: float) -> float:\n", + " return ((end - start) / start) * 100\n", + "\n", + "\n", + "async def main() -> None:\n", + " model_client = OpenAIChatCompletionClient(model=\"gpt-4o\")\n", + "\n", + " # Get a tracer with the default tracer provider.\n", + " tracer = trace.get_tracer(\"tracing-autogen-agentchat\")\n", + "\n", + " # Use the tracer to create a span for the main function.\n", + " with tracer.start_as_current_span(\"run_team\"):\n", + " planning_agent = AssistantAgent(\n", + " \"PlanningAgent\",\n", + " description=\"An agent for planning tasks, this agent should be the first to engage when given a new task.\",\n", + " model_client=model_client,\n", + " system_message=\"\"\"\n", + " You are a planning agent.\n", + " Your job is to break down complex tasks into smaller, manageable subtasks.\n", + " Your team members are:\n", + " WebSearchAgent: Searches for information\n", + " DataAnalystAgent: Performs calculations\n", + "\n", + " You only plan and delegate tasks - you do not execute them yourself.\n", + "\n", + " When assigning tasks, use this format:\n", + " 1. : \n", + "\n", + " After all tasks are complete, summarize the findings and end with \"TERMINATE\".\n", + " \"\"\",\n", + " )\n", + "\n", + " web_search_agent = AssistantAgent(\n", + " \"WebSearchAgent\",\n", + " description=\"An agent for searching information on the web.\",\n", + " tools=[search_web_tool],\n", + " model_client=model_client,\n", + " system_message=\"\"\"\n", + " You are a web search agent.\n", + " Your only tool is search_tool - use it to find information.\n", + " You make only one search call at a time.\n", + " Once you have the results, you never do calculations based on them.\n", + " \"\"\",\n", + " )\n", + "\n", + " data_analyst_agent = AssistantAgent(\n", + " \"DataAnalystAgent\",\n", + " description=\"An agent for performing calculations.\",\n", + " model_client=model_client,\n", + " tools=[percentage_change_tool],\n", + " system_message=\"\"\"\n", + " You are a data analyst.\n", + " Given the tasks you have been assigned, you should analyze the data and provide results using the tools provided.\n", + " If you have not seen the data, ask for it.\n", + " \"\"\",\n", + " )\n", + "\n", + " text_mention_termination = TextMentionTermination(\"TERMINATE\")\n", + " max_messages_termination = MaxMessageTermination(max_messages=25)\n", + " termination = text_mention_termination | max_messages_termination\n", + "\n", + " selector_prompt = \"\"\"Select an agent to perform task.\n", + "\n", + " {roles}\n", + "\n", + " Current conversation context:\n", + " {history}\n", + "\n", + " Read the above conversation, then select an agent from {participants} to perform the next task.\n", + " Make sure the planner agent has assigned tasks before other agents start working.\n", + " Only select one agent.\n", + " \"\"\"\n", + "\n", + " task = \"Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\"\n", + "\n", + " runtime = SingleThreadedAgentRuntime(\n", + " tracer_provider=trace.NoOpTracerProvider(), # Disable telemetry for runtime.\n", + " )\n", + " runtime.start()\n", + "\n", + " team = SelectorGroupChat(\n", + " [planning_agent, web_search_agent, data_analyst_agent],\n", + " model_client=model_client,\n", + " termination_condition=termination,\n", + " selector_prompt=selector_prompt,\n", + " allow_repeated_speaker=True,\n", + " runtime=runtime,\n", + " )\n", + " await Console(team.run_stream(task=task))\n", + "\n", + " await runtime.stop()\n", + "\n", + " await model_client.close()\n", + "\n", + "\n", + "# asyncio.run(main())" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- TextMessage (user) ----------\n", + "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n", + "---------- TextMessage (PlanningAgent) ----------\n", + "To find the information requested, we need to follow these steps:\n", + "\n", + "1. Identify the Miami Heat player with the highest points during the 2006-2007 season.\n", + "2. Get the total rebounds for that player in both the 2007-2008 and 2008-2009 seasons.\n", + "3. Calculate the percentage change in total rebounds between these two seasons.\n", + "\n", + "Here are the tasks assigned to achieve this:\n", + "\n", + "1. WebSearchAgent: Find the Miami Heat player with the highest points during the 2006-2007 season.\n", + "2. WebSearchAgent: After identifying the player, find the total rebounds for that player in the 2007-2008 and 2008-2009 seasons.\n", + "3. DataAnalystAgent: Calculate the percentage change in the player's total rebounds between the 2007-2008 and 2008-2009 seasons.\n", + "---------- ToolCallRequestEvent (WebSearchAgent) ----------\n", + "[FunctionCall(id='call_hS8yod9l6CYUllDveUffp58e', arguments='{\"query\":\"Miami Heat leading scorer 2006-2007 season\"}', name='search_web_tool')]\n", + "---------- ToolCallExecutionEvent (WebSearchAgent) ----------\n", + "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', name='search_web_tool', call_id='call_hS8yod9l6CYUllDveUffp58e', is_error=False)]\n", + "---------- ToolCallSummaryMessage (WebSearchAgent) ----------\n", + "Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", + " Udonis Haslem: 844 points\n", + " Dwayne Wade: 1397 points\n", + " James Posey: 550 points\n", + " ...\n", + " \n", + "---------- ToolCallRequestEvent (WebSearchAgent) ----------\n", + "[FunctionCall(id='call_bUJxtpxUXFSxECDogye9WL0g', arguments='{\"query\":\"Dwyane Wade total rebounds in 2007-2008 season\"}', name='search_web_tool')]\n", + "---------- ToolCallExecutionEvent (WebSearchAgent) ----------\n", + "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', name='search_web_tool', call_id='call_bUJxtpxUXFSxECDogye9WL0g', is_error=False)]\n", + "---------- ToolCallSummaryMessage (WebSearchAgent) ----------\n", + "The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n", + "---------- ToolCallRequestEvent (WebSearchAgent) ----------\n", + "[FunctionCall(id='call_pgYNSDhhyodtteot56FRktxp', arguments='{\"query\":\"Dwyane Wade total rebounds in 2008-2009 season\"}', name='search_web_tool')]\n", + "---------- ToolCallExecutionEvent (WebSearchAgent) ----------\n", + "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', name='search_web_tool', call_id='call_pgYNSDhhyodtteot56FRktxp', is_error=False)]\n", + "---------- ToolCallSummaryMessage (WebSearchAgent) ----------\n", + "The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n", + "---------- ToolCallRequestEvent (DataAnalystAgent) ----------\n", + "[FunctionCall(id='call_A89acjYHlNDLzG09rVNJ0J6H', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')]\n", + "---------- ToolCallExecutionEvent (DataAnalystAgent) ----------\n", + "[FunctionExecutionResult(content='85.98130841121495', name='percentage_change_tool', call_id='call_A89acjYHlNDLzG09rVNJ0J6H', is_error=False)]\n", + "---------- ToolCallSummaryMessage (DataAnalystAgent) ----------\n", + "85.98130841121495\n", + "---------- TextMessage (PlanningAgent) ----------\n", + "The Miami Heat player with the highest points during the 2006-2007 season was Dwyane Wade, who scored 1,397 points. \n", + "\n", + "The total rebounds for Dwyane Wade in the 2007-2008 season were 214, and in the 2008-2009 season, they were 398.\n", + "\n", + "The percentage change in his total rebounds between these two seasons is approximately 86.0%.\n", + "\n", + "TERMINATE\n" + ] + } + ], + "source": [ + "await main()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can then use the Jaeger UI to view the traces collected from the application run above. \n", + "\n", + "![Jaeger UI](jaeger.png)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.7" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/.gitignore b/python/docs/src/user-guide/agentchat-user-guide/tutorial/.gitignore similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/.gitignore rename to python/docs/src/user-guide/agentchat-user-guide/tutorial/.gitignore diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/agentchat-team.svg b/python/docs/src/user-guide/agentchat-user-guide/tutorial/agentchat-team.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/agentchat-team.svg rename to python/docs/src/user-guide/agentchat-user-guide/tutorial/agentchat-team.svg diff --git a/python/docs/src/user-guide/agentchat-user-guide/tutorial/agents.ipynb b/python/docs/src/user-guide/agentchat-user-guide/tutorial/agents.ipynb new file mode 100644 index 000000000000..efd530e1b80a --- /dev/null +++ b/python/docs/src/user-guide/agentchat-user-guide/tutorial/agents.ipynb @@ -0,0 +1,769 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Agents\n", + "\n", + "AutoGen AgentChat provides a set of preset Agents, each with variations in how an agent might respond to messages.\n", + "All agents share the following attributes and methods:\n", + "\n", + "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.name`: The unique name of the agent.\n", + "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.description`: The description of the agent in text.\n", + "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.run`: The method that runs the agent given a task as a string or a list of messages, and returns a {py:class}`~autogen_agentchat.base.TaskResult`. **Agents are expected to be stateful and this method is expected to be called with new messages, not complete history**.\n", + "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.run_stream`: Same as {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run` but returns an iterator of messages that subclass {py:class}`~autogen_agentchat.messages.BaseAgentEvent` or {py:class}`~autogen_agentchat.messages.BaseChatMessage` followed by a {py:class}`~autogen_agentchat.base.TaskResult` as the last item.\n", + "\n", + "See {py:mod}`autogen_agentchat.messages` for more information on AgentChat message types." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Assistant Agent\n", + "\n", + "{py:class}`~autogen_agentchat.agents.AssistantAgent` is a built-in agent that\n", + "uses a language model and has the ability to use tools.\n", + "\n", + "```{warning}\n", + "{py:class}`~autogen_agentchat.agents.AssistantAgent` is a \"kitchen sink\" agent\n", + "for prototyping and educational purpose -- it is very general.\n", + "Make sure you read the documentation and implementation to understand the design choices.\n", + "Once you fully understand the design, you may want to implement your own agent.\n", + "See [Custom Agent](../custom-agents.ipynb).\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.messages import StructuredMessage\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "# Define a tool that searches the web for information.\n", + "# For simplicity, we will use a mock function here that returns a static string.\n", + "async def web_search(query: str) -> str:\n", + " \"\"\"Find information on the web\"\"\"\n", + " return \"AutoGen is a programming framework for building multi-agent applications.\"\n", + "\n", + "\n", + "# Create an agent that uses the OpenAI GPT-4o model.\n", + "model_client = OpenAIChatCompletionClient(\n", + " model=\"gpt-4.1-nano\",\n", + " # api_key=\"YOUR_API_KEY\",\n", + ")\n", + "agent = AssistantAgent(\n", + " name=\"assistant\",\n", + " model_client=model_client,\n", + " tools=[web_search],\n", + " system_message=\"Use tools to solve tasks.\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Getting Result\n", + "\n", + "We can use the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run` method to get the agent run on a given task." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[TextMessage(source='user', models_usage=None, metadata={}, content='Find information on AutoGen', type='TextMessage'), ToolCallRequestEvent(source='assistant', models_usage=RequestUsage(prompt_tokens=61, completion_tokens=16), metadata={}, content=[FunctionCall(id='call_703i17OLXfztkuioUbkESnea', arguments='{\"query\":\"AutoGen\"}', name='web_search')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='assistant', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='AutoGen is a programming framework for building multi-agent applications.', name='web_search', call_id='call_703i17OLXfztkuioUbkESnea', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='assistant', models_usage=None, metadata={}, content='AutoGen is a programming framework for building multi-agent applications.', type='ToolCallSummaryMessage')]\n" + ] + } + ], + "source": [ + "# Use asyncio.run(agent.run(...)) when running in a script.\n", + "result = await agent.run(task=\"Find information on AutoGen\")\n", + "print(result.messages)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The call to the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run` method\n", + "returns a {py:class}`~autogen_agentchat.base.TaskResult`\n", + "with the list of messages in the {py:attr}`~autogen_agentchat.base.TaskResult.messages` attribute,\n", + "which stores the agent's \"thought process\" as well as the final response.\n", + "\n", + "```{note}\n", + "It is important to note that {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run`\n", + "will update the internal state of the agent -- it will add the messages to the agent's\n", + "message history. You can also call {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run`\n", + "without a task to get the agent to generate responses given its current state.\n", + "```\n", + "\n", + "```{note}\n", + "Unlike in v0.2 AgentChat, the tools are executed by the same agent directly within\n", + "the same call to {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run`.\n", + "By default, the agent will return the result of the tool call as the final response.\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Multi-Modal Input\n", + "\n", + "The {py:class}`~autogen_agentchat.agents.AssistantAgent` can handle multi-modal input\n", + "by providing the input as a {py:class}`~autogen_agentchat.messages.MultiModalMessage`." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from io import BytesIO\n", + "\n", + "import PIL\n", + "import requests\n", + "from autogen_agentchat.messages import MultiModalMessage\n", + "from autogen_core import Image\n", + "\n", + "# Create a multi-modal message with random image and text.\n", + "pil_image = PIL.Image.open(BytesIO(requests.get(\"https://picsum.photos/300/200\").content))\n", + "img = Image(pil_image)\n", + "multi_modal_message = MultiModalMessage(content=[\"Can you describe the content of this image?\", img], source=\"user\")\n", + "img" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "The image depicts a scenic mountain landscape under a clear blue sky. There are several rugged mountain peaks in the background, with some clouds scattered across the sky. In the valley below, there is a body of water, possibly a lake or river, surrounded by greenery. The overall scene conveys a sense of natural beauty and tranquility.\n" + ] + } + ], + "source": [ + "# Use asyncio.run(...) when running in a script.\n", + "result = await agent.run(task=multi_modal_message)\n", + "print(result.messages[-1].content) # type: ignore" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Streaming Messages\n", + "\n", + "We can also stream each message as it is generated by the agent by using the\n", + "{py:meth}`~autogen_agentchat.agents.BaseChatAgent.run_stream` method,\n", + "and use {py:class}`~autogen_agentchat.ui.Console` to print the messages\n", + "as they appear to the console." + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- TextMessage (user) ----------\n", + "Find information on AutoGen\n", + "---------- ToolCallRequestEvent (assistant) ----------\n", + "[FunctionCall(id='call_HOTRhOzXCBm0zSqZCFbHD7YP', arguments='{\"query\":\"AutoGen\"}', name='web_search')]\n", + "[Prompt tokens: 61, Completion tokens: 16]\n", + "---------- ToolCallExecutionEvent (assistant) ----------\n", + "[FunctionExecutionResult(content='AutoGen is a programming framework for building multi-agent applications.', name='web_search', call_id='call_HOTRhOzXCBm0zSqZCFbHD7YP', is_error=False)]\n", + "---------- ToolCallSummaryMessage (assistant) ----------\n", + "AutoGen is a programming framework for building multi-agent applications.\n", + "---------- Summary ----------\n", + "Number of messages: 4\n", + "Finish reason: None\n", + "Total prompt tokens: 61\n", + "Total completion tokens: 16\n", + "Duration: 0.52 seconds\n" + ] + } + ], + "source": [ + "async def assistant_run_stream() -> None:\n", + " # Option 1: read each message from the stream (as shown in the previous example).\n", + " # async for message in agent.run_stream(task=\"Find information on AutoGen\"):\n", + " # print(message)\n", + "\n", + " # Option 2: use Console to print all messages as they appear.\n", + " await Console(\n", + " agent.run_stream(task=\"Find information on AutoGen\"),\n", + " output_stats=True, # Enable stats printing.\n", + " )\n", + "\n", + "\n", + "# Use asyncio.run(assistant_run_stream()) when running in a script.\n", + "await assistant_run_stream()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run_stream` method\n", + "returns an asynchronous generator that yields each message generated by the agent,\n", + "followed by a {py:class}`~autogen_agentchat.base.TaskResult` as the last item.\n", + "\n", + "From the messages, you can observe that the assistant agent utilized the `web_search` tool to\n", + "gather information and responded based on the search results." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Tools and Workbench\n", + "\n", + "Large Language Models (LLMs) are typically limited to generating text or code responses. \n", + "However, many complex tasks benefit from the ability to use external tools that perform specific actions,\n", + "such as fetching data from APIs or databases.\n", + "\n", + "To address this limitation, modern LLMs can now accept a list of available tool schemas \n", + "(descriptions of tools and their arguments) and generate a tool call message. \n", + "This capability is known as **Tool Calling** or **Function Calling** and \n", + "is becoming a popular pattern in building intelligent agent-based applications.\n", + "Refer to the documentation from [OpenAI](https://platform.openai.com/docs/guides/function-calling) \n", + "and [Anthropic](https://docs.anthropic.com/en/docs/build-with-claude/tool-use) for more information about tool calling in LLMs.\n", + "\n", + "In AgentChat, the {py:class}`~autogen_agentchat.agents.AssistantAgent` can use tools to perform specific actions.\n", + "The `web_search` tool is one such tool that allows the assistant agent to search the web for information.\n", + "A single custom tool can be a Python function or a subclass of the {py:class}`~autogen_core.tools.BaseTool`.\n", + "\n", + "On the other hand, a {py:class}`~autogen_core.tools.Workbench` is a collection of tools that share state and resources.\n", + "\n", + "```{note}\n", + "For how to use model clients directly with tools and workbench, refer to the [Tools](../../core-user-guide/components/tools.ipynb)\n", + "and [Workbench](../../core-user-guide/components/workbench.ipynb) sections\n", + "in the Core User Guide.\n", + "```\n", + "\n", + "By default, when {py:class}`~autogen_agentchat.agents.AssistantAgent` executes a tool,\n", + "it will return the tool's output as a string in {py:class}`~autogen_agentchat.messages.ToolCallSummaryMessage` in its response.\n", + "If your tool does not return a well-formed string in natural language, you\n", + "can add a reflection step to have the model summarize the tool's output,\n", + "by setting the `reflect_on_tool_use=True` parameter in the {py:class}`~autogen_agentchat.agents.AssistantAgent` constructor.\n", + "\n", + "### Built-in Tools and Workbench\n", + "\n", + "AutoGen Extension provides a set of built-in tools that can be used with the Assistant Agent.\n", + "Head over to the [API documentation](../../../reference/index.md) for all the available tools\n", + "under the `autogen_ext.tools` namespace. For example, you can find the following tools:\n", + "\n", + "- {py:mod}`~autogen_ext.tools.graphrag`: Tools for using GraphRAG index.\n", + "- {py:mod}`~autogen_ext.tools.http`: Tools for making HTTP requests.\n", + "- {py:mod}`~autogen_ext.tools.langchain`: Adaptor for using LangChain tools.\n", + "- {py:mod}`~autogen_ext.tools.mcp`: Tools and workbench for using Model Chat Protocol (MCP) servers." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Function Tool\n", + "\n", + "The {py:class}`~autogen_agentchat.agents.AssistantAgent` automatically\n", + "converts a Python function into a {py:class}`~autogen_core.tools.FunctionTool`\n", + "which can be used as a tool by the agent and automatically generates the tool schema\n", + "from the function signature and docstring.\n", + "\n", + "The `web_search_func` tool is an example of a function tool.\n", + "The schema is automatically generated." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'name': 'web_search_func',\n", + " 'description': 'Find information on the web',\n", + " 'parameters': {'type': 'object',\n", + " 'properties': {'query': {'description': 'query',\n", + " 'title': 'Query',\n", + " 'type': 'string'}},\n", + " 'required': ['query'],\n", + " 'additionalProperties': False},\n", + " 'strict': False}" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from autogen_core.tools import FunctionTool\n", + "\n", + "\n", + "# Define a tool using a Python function.\n", + "async def web_search_func(query: str) -> str:\n", + " \"\"\"Find information on the web\"\"\"\n", + " return \"AutoGen is a programming framework for building multi-agent applications.\"\n", + "\n", + "\n", + "# This step is automatically performed inside the AssistantAgent if the tool is a Python function.\n", + "web_search_function_tool = FunctionTool(web_search_func, description=\"Find information on the web\")\n", + "# The schema is provided to the model during AssistantAgent's on_messages call.\n", + "web_search_function_tool.schema" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Model Context Protocol (MCP) Workbench\n", + "\n", + "The {py:class}`~autogen_agentchat.agents.AssistantAgent` can also use tools that are\n", + "served from a Model Context Protocol (MCP) server\n", + "using {py:func}`~autogen_ext.tools.mcp.McpWorkbench`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Seattle is a major city located in the state of Washington, United States. It was founded on November 13, 1851, and incorporated as a town on January 14, 1865, and later as a city on December 2, 1869. The city is named after Chief Seattle. It covers an area of approximately 142 square miles, with a population of around 737,000 as of the 2020 Census, and an estimated 755,078 residents in 2023. Seattle is known by nicknames such as The Emerald City, Jet City, and Rain City, and has mottos including The City of Flowers and The City of Goodwill. The city operates under a mayor–council government system, with Bruce Harrell serving as mayor. Key landmarks include the Space Needle, Pike Place Market, Amazon Spheres, and the Seattle Great Wheel. It is situated on the U.S. West Coast, with a diverse urban and metropolitan area that extends to a population of over 4 million in the greater metropolitan region.\n" + ] + } + ], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.messages import TextMessage\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams\n", + "\n", + "# Get the fetch tool from mcp-server-fetch.\n", + "fetch_mcp_server = StdioServerParams(command=\"uvx\", args=[\"mcp-server-fetch\"])\n", + "\n", + "# Create an MCP workbench which provides a session to the mcp server.\n", + "async with McpWorkbench(fetch_mcp_server) as workbench: # type: ignore\n", + " # Create an agent that can use the fetch tool.\n", + " model_client = OpenAIChatCompletionClient(model=\"gpt-4.1-nano\")\n", + " fetch_agent = AssistantAgent(\n", + " name=\"fetcher\", model_client=model_client, workbench=workbench, reflect_on_tool_use=True\n", + " )\n", + "\n", + " # Let the agent fetch the content of a URL and summarize it.\n", + " result = await fetch_agent.run(task=\"Summarize the content of https://en.wikipedia.org/wiki/Seattle\")\n", + " assert isinstance(result.messages[-1], TextMessage)\n", + " print(result.messages[-1].content)\n", + "\n", + " # Close the connection to the model client.\n", + " await model_client.close()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Agent as a Tool\n", + "\n", + "Any {py:class}`~autogen_agentchat.agents.BaseChatAgent` can be used as a tool\n", + "by wrapping it in a {py:class}`~autogen_agentchat.tools.AgentTool`.\n", + "This allows for a dynamic, model-driven multi-agent workflow where\n", + "the agent can call other agents as tools to solve tasks." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Parallel Tool Calls\n", + "\n", + "Some models support parallel tool calls, which can be useful for tasks that require multiple tools to be called simultaneously.\n", + "By default, if the model client produces multiple tool calls, {py:class}`~autogen_agentchat.agents.AssistantAgent`\n", + "will call the tools in parallel.\n", + "\n", + "You may want to disable parallel tool calls when the tools have side effects that may interfere with each other, or,\n", + "when agent behavior needs to be consistent across different models.\n", + "This should be done at the model client level.\n", + "\n", + "```{important}\n", + "When using {py:class}`~autogen_agentchat.tools.AgentTool` or {py:class}`~autogen_agentchat.tools.TeamTool`,\n", + "you **must** disable parallel tool calls to avoid concurrency issues.\n", + "These tools cannot run concurrently as agents and teams maintain internal state\n", + "that would conflict with parallel execution.\n", + "```\n", + "\n", + "For {py:class}`~autogen_ext.models.openai.OpenAIChatCompletionClient` and {py:class}`~autogen_ext.models.openai.AzureOpenAIChatCompletionClient`,\n", + "set `parallel_tool_calls=False` to disable parallel tool calls." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "model_client_no_parallel_tool_call = OpenAIChatCompletionClient(\n", + " model=\"gpt-4o\",\n", + " parallel_tool_calls=False, # type: ignore\n", + ")\n", + "agent_no_parallel_tool_call = AssistantAgent(\n", + " name=\"assistant\",\n", + " model_client=model_client_no_parallel_tool_call,\n", + " tools=[web_search],\n", + " system_message=\"Use tools to solve tasks.\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Tool Iterations\n", + "\n", + "One model call followed by one tool call or parallel tool calls\n", + "is a single tool iteration.\n", + "By default, the {py:class}`~autogen_agentchat.agents.AssistantAgent` will\n", + "execute at most one iteration.\n", + "\n", + "The agent can be configured to execute multiple iterations until the model\n", + "stops generating tool calls or the maximum number of iterations is reached.\n", + "You can control the maximum number of iterations by setting the `max_tool_iterations` parameter\n", + "in the {py:class}`~autogen_agentchat.agents.AssistantAgent` constructor." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "agent_loop = AssistantAgent(\n", + " name=\"assistant_loop\",\n", + " model_client=model_client_no_parallel_tool_call,\n", + " tools=[web_search],\n", + " system_message=\"Use tools to solve tasks.\",\n", + " max_tool_iterations=10, # At most 10 iterations of tool calls before stopping the loop.\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Structured Output\n", + "\n", + "Structured output allows models to return structured JSON text with pre-defined schema\n", + "provided by the application. Different from JSON-mode, the schema can be provided\n", + "as a [Pydantic BaseModel](https://docs.pydantic.dev/latest/concepts/models/)\n", + "class, which can also be used to validate the output.\n", + "\n", + "Once you specify the base model class in the `output_content_type` parameter\n", + "of the {py:class}`~autogen_agentchat.agents.AssistantAgent` constructor,\n", + "the agent will respond with a {py:class}`~autogen_agentchat.messages.StructuredMessage`\n", + "whose `content`'s type is the type of the base model class.\n", + "\n", + "This way, you can integrate agent's response directly into your application\n", + "and use the model's output as a structured object.\n", + "\n", + "```{note}\n", + "When the `output_content_type` is set, it by default requires the agent to reflect on the tool use\n", + "and return the a structured output message based on the tool call result.\n", + "You can disable this behavior by setting `reflect_on_tool_use=False` explictly.\n", + "```\n", + "\n", + "Structured output is also useful for incorporating Chain-of-Thought\n", + "reasoning in the agent's responses.\n", + "See the example below for how to use structured output with the assistant agent." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "I am happy.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- assistant ----------\n", + "{\n", + " \"thoughts\": \"The user explicitly states they are happy.\",\n", + " \"response\": \"happy\"\n", + "}\n", + "Thought: The user explicitly states they are happy.\n", + "Response: happy\n" + ] + } + ], + "source": [ + "from typing import Literal\n", + "\n", + "from pydantic import BaseModel\n", + "\n", + "\n", + "# The response format for the agent as a Pydantic base model.\n", + "class AgentResponse(BaseModel):\n", + " thoughts: str\n", + " response: Literal[\"happy\", \"sad\", \"neutral\"]\n", + "\n", + "\n", + "# Create an agent that uses the OpenAI GPT-4o model.\n", + "model_client = OpenAIChatCompletionClient(model=\"gpt-4o\")\n", + "agent = AssistantAgent(\n", + " \"assistant\",\n", + " model_client=model_client,\n", + " system_message=\"Categorize the input as happy, sad, or neutral following the JSON format.\",\n", + " # Define the output content type of the agent.\n", + " output_content_type=AgentResponse,\n", + ")\n", + "\n", + "result = await Console(agent.run_stream(task=\"I am happy.\"))\n", + "\n", + "# Check the last message in the result, validate its type, and print the thoughts and response.\n", + "assert isinstance(result.messages[-1], StructuredMessage)\n", + "assert isinstance(result.messages[-1].content, AgentResponse)\n", + "print(\"Thought: \", result.messages[-1].content.thoughts)\n", + "print(\"Response: \", result.messages[-1].content.response)\n", + "await model_client.close()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Streaming Tokens\n", + "\n", + "You can stream the tokens generated by the model client by setting `model_client_stream=True`.\n", + "This will cause the agent to yield {py:class}`~autogen_agentchat.messages.ModelClientStreamingChunkEvent` messages\n", + "in {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run_stream`.\n", + "\n", + "The underlying model API must support streaming tokens for this to work.\n", + "Please check with your model provider to see if this is supported." + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "source='user' models_usage=None metadata={} content='Name two cities in South America' type='TextMessage'\n", + "source='assistant' models_usage=None metadata={} content='Two' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' cities' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' in' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' South' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' America' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' are' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' Buenos' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' Aires' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' in' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' Argentina' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' and' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' São' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' Paulo' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' in' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content=' Brazil' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=None metadata={} content='.' type='ModelClientStreamingChunkEvent'\n", + "source='assistant' models_usage=RequestUsage(prompt_tokens=0, completion_tokens=0) metadata={} content='Two cities in South America are Buenos Aires in Argentina and São Paulo in Brazil.' type='TextMessage'\n", + "messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Name two cities in South America', type='TextMessage'), TextMessage(source='assistant', models_usage=RequestUsage(prompt_tokens=0, completion_tokens=0), metadata={}, content='Two cities in South America are Buenos Aires in Argentina and São Paulo in Brazil.', type='TextMessage')] stop_reason=None\n" + ] + } + ], + "source": [ + "model_client = OpenAIChatCompletionClient(model=\"gpt-4o\")\n", + "\n", + "streaming_assistant = AssistantAgent(\n", + " name=\"assistant\",\n", + " model_client=model_client,\n", + " system_message=\"You are a helpful assistant.\",\n", + " model_client_stream=True, # Enable streaming tokens.\n", + ")\n", + "\n", + "# Use an async function and asyncio.run() in a script.\n", + "async for message in streaming_assistant.run_stream(task=\"Name two cities in South America\"): # type: ignore\n", + " print(message)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can see the streaming chunks in the output above.\n", + "The chunks are generated by the model client and are yielded by the agent as they are received.\n", + "The final response, the concatenation of all the chunks, is yielded right after the last chunk." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Model Context\n", + "\n", + "{py:class}`~autogen_agentchat.agents.AssistantAgent` has a `model_context`\n", + "parameter that can be used to pass in a {py:class}`~autogen_core.model_context.ChatCompletionContext`\n", + "object. This allows the agent to use different model contexts, such as\n", + "{py:class}`~autogen_core.model_context.BufferedChatCompletionContext` to\n", + "limit the context sent to the model.\n", + "\n", + "By default, {py:class}`~autogen_agentchat.agents.AssistantAgent` uses\n", + "the {py:class}`~autogen_core.model_context.UnboundedChatCompletionContext`\n", + "which sends the full conversation history to the model. To limit the context\n", + "to the last `n` messages, you can use the {py:class}`~autogen_core.model_context.BufferedChatCompletionContext`.\n", + "To limit the context by token count, you can use the\n", + "{py:class}`~autogen_core.model_context.TokenLimitedChatCompletionContext`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from autogen_core.model_context import BufferedChatCompletionContext\n", + "\n", + "# Create an agent that uses only the last 5 messages in the context to generate responses.\n", + "agent = AssistantAgent(\n", + " name=\"assistant\",\n", + " model_client=model_client,\n", + " tools=[web_search],\n", + " system_message=\"Use tools to solve tasks.\",\n", + " model_context=BufferedChatCompletionContext(buffer_size=5), # Only use the last 5 messages in the context.\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Other Preset Agents\n", + "\n", + "The following preset agents are available:\n", + "\n", + "- {py:class}`~autogen_agentchat.agents.UserProxyAgent`: An agent that takes user input returns it as responses.\n", + "- {py:class}`~autogen_agentchat.agents.CodeExecutorAgent`: An agent that can execute code.\n", + "- {py:class}`~autogen_ext.agents.openai.OpenAIAssistantAgent`: An agent that is backed by an OpenAI Assistant, with ability to use custom tools.\n", + "- {py:class}`~autogen_ext.agents.web_surfer.MultimodalWebSurfer`: A multi-modal agent that can search the web and visit web pages for information.\n", + "- {py:class}`~autogen_ext.agents.file_surfer.FileSurfer`: An agent that can search and browse local files for information.\n", + "- {py:class}`~autogen_ext.agents.video_surfer.VideoSurfer`: An agent that can watch videos for information." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Next Step\n", + "\n", + "Having explored the usage of the {py:class}`~autogen_agentchat.agents.AssistantAgent`, we can now proceed to the next section to learn about the teams feature in AgentChat.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "python", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/human-in-the-loop-termination.svg b/python/docs/src/user-guide/agentchat-user-guide/tutorial/human-in-the-loop-termination.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/human-in-the-loop-termination.svg rename to python/docs/src/user-guide/agentchat-user-guide/tutorial/human-in-the-loop-termination.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/human-in-the-loop-user-proxy.svg b/python/docs/src/user-guide/agentchat-user-guide/tutorial/human-in-the-loop-user-proxy.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/human-in-the-loop-user-proxy.svg rename to python/docs/src/user-guide/agentchat-user-guide/tutorial/human-in-the-loop-user-proxy.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/human-in-the-loop.ipynb b/python/docs/src/user-guide/agentchat-user-guide/tutorial/human-in-the-loop.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/human-in-the-loop.ipynb rename to python/docs/src/user-guide/agentchat-user-guide/tutorial/human-in-the-loop.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/index.md b/python/docs/src/user-guide/agentchat-user-guide/tutorial/index.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/index.md rename to python/docs/src/user-guide/agentchat-user-guide/tutorial/index.md diff --git a/python/docs/src/user-guide/agentchat-user-guide/tutorial/messages.ipynb b/python/docs/src/user-guide/agentchat-user-guide/tutorial/messages.ipynb new file mode 100644 index 000000000000..24b3cc61eb18 --- /dev/null +++ b/python/docs/src/user-guide/agentchat-user-guide/tutorial/messages.ipynb @@ -0,0 +1,137 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Messages" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In AutoGen AgentChat, _messages_ facilitate communication and information exchange with other agents, orchestrators, and applications. AgentChat supports various message types, each designed for specific purposes." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Types of Messages\n", + "\n", + "At a high level, messages in AgentChat can be categorized into two types: agent-agent messages and an agent's internal events and messages.\n", + "\n", + "### Agent-Agent Messages\n", + "AgentChat supports many message types for agent-to-agent communication. They belong to subclasses of the base class {py:class}`~autogen_agentchat.messages.BaseChatMessage`. Concrete subclasses covers basic text and multimodal communication, such as {py:class}`~autogen_agentchat.messages.TextMessage` and {py:class}`~autogen_agentchat.messages.MultiModalMessage`.\n", + "\n", + "For example, the following code snippet demonstrates how to create a text message, which accepts a string content and a string source:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from autogen_agentchat.messages import TextMessage\n", + "\n", + "text_message = TextMessage(content=\"Hello, world!\", source=\"User\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Similarly, the following code snippet demonstrates how to create a multimodal message, which accepts\n", + "a list of strings or {py:class}`~autogen_core.Image` objects:" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from io import BytesIO\n", + "\n", + "import requests\n", + "from autogen_agentchat.messages import MultiModalMessage\n", + "from autogen_core import Image as AGImage\n", + "from PIL import Image\n", + "\n", + "pil_image = Image.open(BytesIO(requests.get(\"https://picsum.photos/300/200\").content))\n", + "img = AGImage(pil_image)\n", + "multi_modal_message = MultiModalMessage(content=[\"Can you describe the content of this image?\", img], source=\"User\")\n", + "img" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The {py:class}`~autogen_agentchat.messages.TextMessage` and {py:class}`~autogen_agentchat.messages.MultiModalMessage` we have created can be passed to agents directly via the {py:class}`~autogen_agentchat.base.ChatAgent.on_messages` method, or as tasks given to a team {py:meth}`~autogen_agentchat.teams.BaseGroupChat.run` method. Messages are also used in the responses of an agent. We will explain these in more detail in [Agents](./agents.ipynb) and [Teams](./teams.ipynb)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Internal Events\n", + "\n", + "AgentChat also supports the concept of `events` - messages that are internal to an agent. These messages are used to communicate events and information on actions _within_ the agent itself, and belong to subclasses of the base class {py:class}`~autogen_agentchat.messages.BaseAgentEvent`.\n", + "\n", + "Examples of these include {py:class}`~autogen_agentchat.messages.ToolCallRequestEvent`, which indicates that a request was made to call a tool, and {py:class}`~autogen_agentchat.messages.ToolCallExecutionEvent`, which contains the results of tool calls.\n", + "\n", + "Typically, events are created by the agent itself and are contained in the {py:attr}`~autogen_agentchat.base.Response.inner_messages` field of the {py:class}`~autogen_agentchat.base.Response` returned from {py:class}`~autogen_agentchat.base.ChatAgent.on_messages`. If you are building a custom agent and have events that you want to communicate to other entities (e.g., a UI), you can include these in the {py:attr}`~autogen_agentchat.base.Response.inner_messages` field of the {py:class}`~autogen_agentchat.base.Response`. We will show examples of this in [Custom Agents](../custom-agents.ipynb).\n", + "\n", + "\n", + "You can read about the full set of messages supported in AgentChat in the {py:mod}`~autogen_agentchat.messages` module. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Custom Message Types\n", + "\n", + "You can create custom message types by subclassing the base class {py:class}`~autogen_agentchat.messages.BaseChatMessage` or {py:class}`~autogen_agentchat.messages.BaseAgentEvent`. This allows you to define your own message formats and behaviors, tailored to your application. Custom message types are useful when you write custom agents." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "agnext", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.9" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/models.ipynb b/python/docs/src/user-guide/agentchat-user-guide/tutorial/models.ipynb similarity index 84% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/models.ipynb rename to python/docs/src/user-guide/agentchat-user-guide/tutorial/models.ipynb index 8674b4195e99..0d1bdf26ba80 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/models.ipynb +++ b/python/docs/src/user-guide/agentchat-user-guide/tutorial/models.ipynb @@ -162,6 +162,7 @@ "metadata": {}, "outputs": [], "source": [ + "from autogen_core.models import UserMessage\n", "from autogen_ext.auth.azure import AzureTokenProvider\n", "from autogen_ext.models.openai import AzureOpenAIChatCompletionClient\n", "from azure.identity import DefaultAzureCredential\n", @@ -247,7 +248,7 @@ "\n", "client = AzureAIChatCompletionClient(\n", " model=\"Phi-4\",\n", - " endpoint=\"https://models.inference.ai.azure.com\",\n", + " endpoint=\"https://models.github.ai/inference\",\n", " # To authenticate with the model you will need to generate a personal access token (PAT) in your GitHub settings.\n", " # Create your PAT token by following instructions here: https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens\n", " credential=AzureKeyCredential(os.environ[\"GITHUB_TOKEN\"]),\n", @@ -373,6 +374,7 @@ "```{note}\n", "While some model providers may offer OpenAI-compatible APIs, they may still have minor differences.\n", "For example, the `finish_reason` field may be different in the response.\n", + "\n", "```" ] }, @@ -403,6 +405,82 @@ "await model_client.close()" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Also, as Gemini adds new models, you may need to define the models capabilities via the model_info field. For example, to use `gemini-2.0-flash-lite` or a similar new model, you can use the following code:\n", + "\n", + "```python \n", + "from autogen_core.models import UserMessage\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "from autogen_core.models import ModelInfo\n", + "\n", + "model_client = OpenAIChatCompletionClient(\n", + " model=\"gemini-2.0-flash-lite\",\n", + " model_info=ModelInfo(vision=True, function_calling=True, json_output=True, family=\"unknown\", structured_output=True)\n", + " # api_key=\"GEMINI_API_KEY\",\n", + ")\n", + "\n", + "response = await model_client.create([UserMessage(content=\"What is the capital of France?\", source=\"user\")])\n", + "print(response)\n", + "await model_client.close()\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Llama API (experimental)\n", + "\n", + "[Llama API](https://llama.developer.meta.com?utm_source=partner-autogen&utm_medium=readme) is the Meta's first party API offering. It currently offers an [OpenAI compatible endpoint](https://llama.developer.meta.com/docs/features/compatibility).\n", + "So you can use the {py:class}`~autogen_ext.models.openai.OpenAIChatCompletionClient` with the Llama API.\n", + "\n", + "This endpoint fully supports the following OpenAI client library features:\n", + "* Chat completions\n", + "* Model selection\n", + "* Temperature/sampling\n", + "* Streaming\n", + "* Image understanding\n", + "* Structured output (JSON mode)\n", + "* Function calling (tools)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "from autogen_core import Image\n", + "from autogen_core.models import UserMessage\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "# Text\n", + "model_client = OpenAIChatCompletionClient(\n", + " model=\"Llama-4-Scout-17B-16E-Instruct-FP8\",\n", + " # api_key=\"LLAMA_API_KEY\"\n", + ")\n", + "\n", + "response = await model_client.create([UserMessage(content=\"Write me a poem\", source=\"user\")])\n", + "print(response)\n", + "await model_client.close()\n", + "\n", + "# Image\n", + "model_client = OpenAIChatCompletionClient(\n", + " model=\"Llama-4-Maverick-17B-128E-Instruct-FP8\",\n", + " # api_key=\"LLAMA_API_KEY\"\n", + ")\n", + "image = Image.from_file(Path(\"test.png\"))\n", + "\n", + "response = await model_client.create([UserMessage(content=[\"What is in this image\", image], source=\"user\")])\n", + "print(response)\n", + "await model_client.close()" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -516,7 +594,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.7" + "version": "3.11.9" } }, "nbformat": 4, diff --git a/python/docs/src/user-guide/agentchat-user-guide/tutorial/state.ipynb b/python/docs/src/user-guide/agentchat-user-guide/tutorial/state.ipynb new file mode 100644 index 000000000000..636841d085ee --- /dev/null +++ b/python/docs/src/user-guide/agentchat-user-guide/tutorial/state.ipynb @@ -0,0 +1,359 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Managing State \n", + "\n", + "So far, we have discussed how to build components in a multi-agent application - agents, teams, termination conditions. In many cases, it is useful to save the state of these components to disk and load them back later. This is particularly useful in a web application where stateless endpoints respond to requests and need to load the state of the application from persistent storage.\n", + "\n", + "In this notebook, we will discuss how to save and load the state of agents, teams, and termination conditions. \n", + " \n", + "\n", + "## Saving and Loading Agents\n", + "\n", + "We can get the state of an agent by calling {py:meth}`~autogen_agentchat.agents.AssistantAgent.save_state` method on \n", + "an {py:class}`~autogen_agentchat.agents.AssistantAgent`. " + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "In Tanganyika's embrace so wide and deep, \n", + "Ancient waters cradle secrets they keep, \n", + "Echoes of time where horizons sleep. \n" + ] + } + ], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.conditions import MaxMessageTermination\n", + "from autogen_agentchat.messages import TextMessage\n", + "from autogen_agentchat.teams import RoundRobinGroupChat\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_core import CancellationToken\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-2024-08-06\")\n", + "\n", + "assistant_agent = AssistantAgent(\n", + " name=\"assistant_agent\",\n", + " system_message=\"You are a helpful assistant\",\n", + " model_client=model_client,\n", + ")\n", + "\n", + "# Use asyncio.run(...) when running in a script.\n", + "response = await assistant_agent.on_messages(\n", + " [TextMessage(content=\"Write a 3 line poem on lake tangayika\", source=\"user\")], CancellationToken()\n", + ")\n", + "print(response.chat_message)\n", + "await model_client.close()" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{'type': 'AssistantAgentState', 'version': '1.0.0', 'llm_messages': [{'content': 'Write a 3 line poem on lake tangayika', 'source': 'user', 'type': 'UserMessage'}, {'content': \"In Tanganyika's embrace so wide and deep, \\nAncient waters cradle secrets they keep, \\nEchoes of time where horizons sleep. \", 'source': 'assistant_agent', 'type': 'AssistantMessage'}]}\n" + ] + } + ], + "source": [ + "agent_state = await assistant_agent.save_state()\n", + "print(agent_state)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "The last line of the poem was: \"Echoes of time where horizons sleep.\"\n" + ] + } + ], + "source": [ + "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-2024-08-06\")\n", + "\n", + "new_assistant_agent = AssistantAgent(\n", + " name=\"assistant_agent\",\n", + " system_message=\"You are a helpful assistant\",\n", + " model_client=model_client,\n", + ")\n", + "await new_assistant_agent.load_state(agent_state)\n", + "\n", + "# Use asyncio.run(...) when running in a script.\n", + "response = await new_assistant_agent.on_messages(\n", + " [TextMessage(content=\"What was the last line of the previous poem you wrote\", source=\"user\")], CancellationToken()\n", + ")\n", + "print(response.chat_message)\n", + "await model_client.close()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "```{note}\n", + "For {py:class}`~autogen_agentchat.agents.AssistantAgent`, its state consists of the model_context.\n", + "If you write your own custom agent, consider overriding the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.save_state` and {py:meth}`~autogen_agentchat.agents.BaseChatAgent.load_state` methods to customize the behavior. The default implementations save and load an empty state.\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Saving and Loading Teams \n", + "\n", + "We can get the state of a team by calling `save_state` method on the team and load it back by calling `load_state` method on the team. \n", + "\n", + "When we call `save_state` on a team, it saves the state of all the agents in the team.\n", + "\n", + "We will begin by creating a simple {py:class}`~autogen_agentchat.teams.RoundRobinGroupChat` team with a single agent and ask it to write a poem. " + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "Write a beautiful poem 3-line about lake tangayika\n", + "---------- assistant_agent ----------\n", + "In Tanganyika's gleam, beneath the azure skies, \n", + "Whispers of ancient waters, in tranquil guise, \n", + "Nature's mirror, where dreams and serenity lie.\n", + "[Prompt tokens: 29, Completion tokens: 34]\n", + "---------- Summary ----------\n", + "Number of messages: 2\n", + "Finish reason: Maximum number of messages 2 reached, current message count: 2\n", + "Total prompt tokens: 29\n", + "Total completion tokens: 34\n", + "Duration: 0.71 seconds\n" + ] + } + ], + "source": [ + "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-2024-08-06\")\n", + "\n", + "# Define a team.\n", + "assistant_agent = AssistantAgent(\n", + " name=\"assistant_agent\",\n", + " system_message=\"You are a helpful assistant\",\n", + " model_client=model_client,\n", + ")\n", + "agent_team = RoundRobinGroupChat([assistant_agent], termination_condition=MaxMessageTermination(max_messages=2))\n", + "\n", + "# Run the team and stream messages to the console.\n", + "stream = agent_team.run_stream(task=\"Write a beautiful poem 3-line about lake tangayika\")\n", + "\n", + "# Use asyncio.run(...) when running in a script.\n", + "await Console(stream)\n", + "\n", + "# Save the state of the agent team.\n", + "team_state = await agent_team.save_state()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If we reset the team (simulating instantiation of the team), and ask the question `What was the last line of the poem you wrote?`, we see that the team is unable to accomplish this as there is no reference to the previous run." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "What was the last line of the poem you wrote?\n", + "---------- assistant_agent ----------\n", + "I'm sorry, but I am unable to recall or access previous interactions, including any specific poem I may have composed in our past conversations. If you like, I can write a new poem for you.\n", + "[Prompt tokens: 28, Completion tokens: 40]\n", + "---------- Summary ----------\n", + "Number of messages: 2\n", + "Finish reason: Maximum number of messages 2 reached, current message count: 2\n", + "Total prompt tokens: 28\n", + "Total completion tokens: 40\n", + "Duration: 0.70 seconds\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What was the last line of the poem you wrote?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=28, completion_tokens=40), content=\"I'm sorry, but I am unable to recall or access previous interactions, including any specific poem I may have composed in our past conversations. If you like, I can write a new poem for you.\", type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "await agent_team.reset()\n", + "stream = agent_team.run_stream(task=\"What was the last line of the poem you wrote?\")\n", + "await Console(stream)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Next, we load the state of the team and ask the same question. We see that the team is able to accurately return the last line of the poem it wrote." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{'type': 'TeamState', 'version': '1.0.0', 'agent_states': {'group_chat_manager/a55364ad-86fd-46ab-9449-dcb5260b1e06': {'type': 'RoundRobinManagerState', 'version': '1.0.0', 'message_thread': [{'source': 'user', 'models_usage': None, 'content': 'Write a beautiful poem 3-line about lake tangayika', 'type': 'TextMessage'}, {'source': 'assistant_agent', 'models_usage': {'prompt_tokens': 29, 'completion_tokens': 34}, 'content': \"In Tanganyika's gleam, beneath the azure skies, \\nWhispers of ancient waters, in tranquil guise, \\nNature's mirror, where dreams and serenity lie.\", 'type': 'TextMessage'}], 'current_turn': 0, 'next_speaker_index': 0}, 'collect_output_messages/a55364ad-86fd-46ab-9449-dcb5260b1e06': {}, 'assistant_agent/a55364ad-86fd-46ab-9449-dcb5260b1e06': {'type': 'ChatAgentContainerState', 'version': '1.0.0', 'agent_state': {'type': 'AssistantAgentState', 'version': '1.0.0', 'llm_messages': [{'content': 'Write a beautiful poem 3-line about lake tangayika', 'source': 'user', 'type': 'UserMessage'}, {'content': \"In Tanganyika's gleam, beneath the azure skies, \\nWhispers of ancient waters, in tranquil guise, \\nNature's mirror, where dreams and serenity lie.\", 'source': 'assistant_agent', 'type': 'AssistantMessage'}]}, 'message_buffer': []}}, 'team_id': 'a55364ad-86fd-46ab-9449-dcb5260b1e06'}\n", + "---------- user ----------\n", + "What was the last line of the poem you wrote?\n", + "---------- assistant_agent ----------\n", + "The last line of the poem I wrote is: \n", + "\"Nature's mirror, where dreams and serenity lie.\"\n", + "[Prompt tokens: 86, Completion tokens: 22]\n", + "---------- Summary ----------\n", + "Number of messages: 2\n", + "Finish reason: Maximum number of messages 2 reached, current message count: 2\n", + "Total prompt tokens: 86\n", + "Total completion tokens: 22\n", + "Duration: 0.96 seconds\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What was the last line of the poem you wrote?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=22), content='The last line of the poem I wrote is: \\n\"Nature\\'s mirror, where dreams and serenity lie.\"', type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "print(team_state)\n", + "\n", + "# Load team state.\n", + "await agent_team.load_state(team_state)\n", + "stream = agent_team.run_stream(task=\"What was the last line of the poem you wrote?\")\n", + "await Console(stream)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Persisting State (File or Database)\n", + "\n", + "In many cases, we may want to persist the state of the team to disk (or a database) and load it back later. State is a dictionary that can be serialized to a file or written to a database." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "What was the last line of the poem you wrote?\n", + "---------- assistant_agent ----------\n", + "The last line of the poem I wrote is: \n", + "\"Nature's mirror, where dreams and serenity lie.\"\n", + "[Prompt tokens: 86, Completion tokens: 22]\n", + "---------- Summary ----------\n", + "Number of messages: 2\n", + "Finish reason: Maximum number of messages 2 reached, current message count: 2\n", + "Total prompt tokens: 86\n", + "Total completion tokens: 22\n", + "Duration: 0.72 seconds\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What was the last line of the poem you wrote?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=22), content='The last line of the poem I wrote is: \\n\"Nature\\'s mirror, where dreams and serenity lie.\"', type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import json\n", + "\n", + "## save state to disk\n", + "\n", + "with open(\"coding/team_state.json\", \"w\") as f:\n", + " json.dump(team_state, f)\n", + "\n", + "## load state from disk\n", + "with open(\"coding/team_state.json\", \"r\") as f:\n", + " team_state = json.load(f)\n", + "\n", + "new_agent_team = RoundRobinGroupChat([assistant_agent], termination_condition=MaxMessageTermination(max_messages=2))\n", + "await new_agent_team.load_state(team_state)\n", + "stream = new_agent_team.run_stream(task=\"What was the last line of the poem you wrote?\")\n", + "await Console(stream)\n", + "await model_client.close()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "agnext", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.9" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/teams.ipynb b/python/docs/src/user-guide/agentchat-user-guide/tutorial/teams.ipynb similarity index 98% rename from python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/teams.ipynb rename to python/docs/src/user-guide/agentchat-user-guide/tutorial/teams.ipynb index e5419c0a891b..9a1bab07b6ce 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/teams.ipynb +++ b/python/docs/src/user-guide/agentchat-user-guide/tutorial/teams.ipynb @@ -16,6 +16,7 @@ "- {py:class}`~autogen_agentchat.teams.RoundRobinGroupChat`: A team that runs a group chat with participants taking turns in a round-robin fashion (covered on this page). [Tutorial](#creating-a-team) \n", "- {py:class}`~autogen_agentchat.teams.SelectorGroupChat`: A team that selects the next speaker using a ChatCompletion model after each message. [Tutorial](../selector-group-chat.ipynb)\n", "- {py:class}`~autogen_agentchat.teams.MagenticOneGroupChat`: A generalist multi-agent system for solving open-ended web and file-based tasks across a variety of domains. [Tutorial](../magentic-one.md) \n", + "- {py:class}`~autogen_agentchat.teams.Swarm`: A team that uses {py:class}`~autogen_agentchat.messages.HandoffMessage` to signal transitions between agents. [Tutorial](../swarm.ipynb)\n", "\n", "```{note}\n", "\n", @@ -589,6 +590,13 @@ "source": [ "## Single-Agent Team\n", "\n", + "```{note}\n", + "Starting with version 0.6.2, you can use {py:class}`~autogen_agentchat.agents.AssistantAgent`\n", + "with `max_tool_iterations` to run the agent with multiple iterations\n", + "of tool calls. So you may not need to use a single-agent team if you just \n", + "want to run the agent in a tool-calling loop.\n", + "```\n", + "\n", "Often, you may want to run a single agent in a team configuration.\n", "This is useful for running the {py:class}`~autogen_agentchat.agents.AssistantAgent` in a loop\n", "until a termination condition is met.\n", @@ -697,7 +705,7 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "python", "language": "python", "name": "python3" }, @@ -711,7 +719,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.3" + "version": "3.12.11" } }, "nbformat": 4, diff --git a/python/docs/src/user-guide/agentchat-user-guide/tutorial/termination.ipynb b/python/docs/src/user-guide/agentchat-user-guide/tutorial/termination.ipynb new file mode 100644 index 000000000000..b12b874043a0 --- /dev/null +++ b/python/docs/src/user-guide/agentchat-user-guide/tutorial/termination.ipynb @@ -0,0 +1,519 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Termination \n", + "\n", + "In the previous section, we explored how to define agents, and organize them into teams that can solve tasks. However, a run can go on forever, and in many cases, we need to know _when_ to stop them. This is the role of the termination condition.\n", + "\n", + "AgentChat supports several termination condition by providing a base {py:class}`~autogen_agentchat.base.TerminationCondition` class and several implementations that inherit from it.\n", + "\n", + "A termination condition is a callable that takes a sequence of {py:class}`~autogen_agentchat.messages.BaseAgentEvent` or {py:class}`~autogen_agentchat.messages.BaseChatMessage` objects **since the last time the condition was called**, and returns a {py:class}`~autogen_agentchat.messages.StopMessage` if the conversation should be terminated, or `None` otherwise.\n", + "Once a termination condition has been reached, it must be reset by calling {py:meth}`~autogen_agentchat.base.TerminationCondition.reset` before it can be used again.\n", + "\n", + "Some important things to note about termination conditions: \n", + "- They are stateful but reset automatically after each run ({py:meth}`~autogen_agentchat.base.TaskRunner.run` or {py:meth}`~autogen_agentchat.base.TaskRunner.run_stream`) is finished.\n", + "- They can be combined using the AND and OR operators.\n", + "\n", + "```{note}\n", + "For group chat teams (i.e., {py:class}`~autogen_agentchat.teams.RoundRobinGroupChat`,\n", + "{py:class}`~autogen_agentchat.teams.SelectorGroupChat`, and {py:class}`~autogen_agentchat.teams.Swarm`),\n", + "the termination condition is called after each agent responds.\n", + "While a response may contain multiple inner messages, the team calls its termination condition just once for all the messages from a single response.\n", + "So the condition is called with the \"delta sequence\" of messages since the last time it was called.\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Built-In Termination Conditions: \n", + "1. {py:class}`~autogen_agentchat.conditions.MaxMessageTermination`: Stops after a specified number of messages have been produced, including both agent and task messages.\n", + "2. {py:class}`~autogen_agentchat.conditions.TextMentionTermination`: Stops when specific text or string is mentioned in a message (e.g., \"TERMINATE\").\n", + "3. {py:class}`~autogen_agentchat.conditions.TokenUsageTermination`: Stops when a certain number of prompt or completion tokens are used. This requires the agents to report token usage in their messages.\n", + "4. {py:class}`~autogen_agentchat.conditions.TimeoutTermination`: Stops after a specified duration in seconds.\n", + "5. {py:class}`~autogen_agentchat.conditions.HandoffTermination`: Stops when a handoff to a specific target is requested. Handoff messages can be used to build patterns such as {py:class}`~autogen_agentchat.teams.Swarm`. This is useful when you want to pause the run and allow application or user to provide input when an agent hands off to them.\n", + "6. {py:class}`~autogen_agentchat.conditions.SourceMatchTermination`: Stops after a specific agent responds.\n", + "7. {py:class}`~autogen_agentchat.conditions.ExternalTermination`: Enables programmatic control of termination from outside the run. This is useful for UI integration (e.g., \"Stop\" buttons in chat interfaces).\n", + "8. {py:class}`~autogen_agentchat.conditions.StopMessageTermination`: Stops when a {py:class}`~autogen_agentchat.messages.StopMessage` is produced by an agent.\n", + "9. {py:class}`~autogen_agentchat.conditions.TextMessageTermination`: Stops when a {py:class}`~autogen_agentchat.messages.TextMessage` is produced by an agent.\n", + "10. {py:class}`~autogen_agentchat.conditions.FunctionCallTermination`: Stops when a {py:class}`~autogen_agentchat.messages.ToolCallExecutionEvent` containing a {py:class}`~autogen_core.models.FunctionExecutionResult` with a matching name is produced by an agent.\n", + "11. {py:class}`~autogen_agentchat.conditions.FunctionalTermination`: Stop when a function expression is evaluated to `True` on the last delta sequence of messages. This is useful for quickly create custom termination conditions that are not covered by the built-in ones." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Basic Usage\n", + "\n", + "To demonstrate the characteristics of termination conditions, we'll create a team consisting of two agents: a primary agent responsible for text generation and a critic agent that reviews and provides feedback on the generated text." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination\n", + "from autogen_agentchat.teams import RoundRobinGroupChat\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "model_client = OpenAIChatCompletionClient(\n", + " model=\"gpt-4o\",\n", + " temperature=1,\n", + " # api_key=\"sk-...\", # Optional if you have an OPENAI_API_KEY env variable set.\n", + ")\n", + "\n", + "# Create the primary agent.\n", + "primary_agent = AssistantAgent(\n", + " \"primary\",\n", + " model_client=model_client,\n", + " system_message=\"You are a helpful AI assistant.\",\n", + ")\n", + "\n", + "# Create the critic agent.\n", + "critic_agent = AssistantAgent(\n", + " \"critic\",\n", + " model_client=model_client,\n", + " system_message=\"Provide constructive feedback for every message. Respond with 'APPROVE' to when your feedbacks are addressed.\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Let's explore how termination conditions automatically reset after each `run` or `run_stream` call, allowing the team to resume its conversation from where it left off." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "Write a unique, Haiku about the weather in Paris\n", + "---------- primary ----------\n", + "Gentle rain whispers, \n", + "Cobblestones glisten softly— \n", + "Paris dreams in gray.\n", + "[Prompt tokens: 30, Completion tokens: 19]\n", + "---------- critic ----------\n", + "The Haiku captures the essence of a rainy day in Paris beautifully, and the imagery is vivid. However, it's important to ensure the use of the traditional 5-7-5 syllable structure for Haikus. Your current Haiku lines are composed of 4-7-5 syllables, which slightly deviates from the form. Consider revising the first line to fit the structure.\n", + "\n", + "For example:\n", + "Soft rain whispers down, \n", + "Cobblestones glisten softly — \n", + "Paris dreams in gray.\n", + "\n", + "This revision maintains the essence of your original lines while adhering to the traditional Haiku structure.\n", + "[Prompt tokens: 70, Completion tokens: 120]\n", + "---------- Summary ----------\n", + "Number of messages: 3\n", + "Finish reason: Maximum number of messages 3 reached, current message count: 3\n", + "Total prompt tokens: 100\n", + "Total completion tokens: 139\n", + "Duration: 3.34 seconds\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Write a unique, Haiku about the weather in Paris'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=30, completion_tokens=19), content='Gentle rain whispers, \\nCobblestones glisten softly— \\nParis dreams in gray.'), TextMessage(source='critic', models_usage=RequestUsage(prompt_tokens=70, completion_tokens=120), content=\"The Haiku captures the essence of a rainy day in Paris beautifully, and the imagery is vivid. However, it's important to ensure the use of the traditional 5-7-5 syllable structure for Haikus. Your current Haiku lines are composed of 4-7-5 syllables, which slightly deviates from the form. Consider revising the first line to fit the structure.\\n\\nFor example:\\nSoft rain whispers down, \\nCobblestones glisten softly — \\nParis dreams in gray.\\n\\nThis revision maintains the essence of your original lines while adhering to the traditional Haiku structure.\")], stop_reason='Maximum number of messages 3 reached, current message count: 3')" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "max_msg_termination = MaxMessageTermination(max_messages=3)\n", + "round_robin_team = RoundRobinGroupChat([primary_agent, critic_agent], termination_condition=max_msg_termination)\n", + "\n", + "# Use asyncio.run(...) if you are running this script as a standalone script.\n", + "await Console(round_robin_team.run_stream(task=\"Write a unique, Haiku about the weather in Paris\"))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The conversation stopped after reaching the maximum message limit. Since the primary agent didn't get to respond to the feedback, let's continue the conversation." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- primary ----------\n", + "Thank you for your feedback. Here is the revised Haiku:\n", + "\n", + "Soft rain whispers down, \n", + "Cobblestones glisten softly — \n", + "Paris dreams in gray.\n", + "[Prompt tokens: 181, Completion tokens: 32]\n", + "---------- critic ----------\n", + "The revised Haiku now follows the traditional 5-7-5 syllable pattern, and it still beautifully captures the atmospheric mood of Paris in the rain. The imagery and flow are both clear and evocative. Well done on making the adjustment! \n", + "\n", + "APPROVE\n", + "[Prompt tokens: 234, Completion tokens: 54]\n", + "---------- primary ----------\n", + "Thank you for your kind words and approval. I'm glad the revision meets your expectations and captures the essence of Paris. If you have any more requests or need further assistance, feel free to ask!\n", + "[Prompt tokens: 279, Completion tokens: 39]\n", + "---------- Summary ----------\n", + "Number of messages: 3\n", + "Finish reason: Maximum number of messages 3 reached, current message count: 3\n", + "Total prompt tokens: 694\n", + "Total completion tokens: 125\n", + "Duration: 6.43 seconds\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=181, completion_tokens=32), content='Thank you for your feedback. Here is the revised Haiku:\\n\\nSoft rain whispers down, \\nCobblestones glisten softly — \\nParis dreams in gray.'), TextMessage(source='critic', models_usage=RequestUsage(prompt_tokens=234, completion_tokens=54), content='The revised Haiku now follows the traditional 5-7-5 syllable pattern, and it still beautifully captures the atmospheric mood of Paris in the rain. The imagery and flow are both clear and evocative. Well done on making the adjustment! \\n\\nAPPROVE'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=279, completion_tokens=39), content=\"Thank you for your kind words and approval. I'm glad the revision meets your expectations and captures the essence of Paris. If you have any more requests or need further assistance, feel free to ask!\")], stop_reason='Maximum number of messages 3 reached, current message count: 3')" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Use asyncio.run(...) if you are running this script as a standalone script.\n", + "await Console(round_robin_team.run_stream())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The team continued from where it left off, allowing the primary agent to respond to the feedback." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Combining Termination Conditions\n", + "\n", + "Let's show how termination conditions can be combined using the AND (`&`) and OR (`|`) operators to create more complex termination logic. For example, we'll create a team that stops either after 10 messages are generated or when the critic agent approves a message.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "Write a unique, Haiku about the weather in Paris\n", + "---------- primary ----------\n", + "Spring breeze gently hums, \n", + "Cherry blossoms in full bloom— \n", + "Paris wakes to life.\n", + "[Prompt tokens: 467, Completion tokens: 19]\n", + "---------- critic ----------\n", + "The Haiku beautifully captures the awakening of Paris in the spring. The imagery of a gentle spring breeze and cherry blossoms in full bloom effectively conveys the rejuvenating feel of the season. The final line, \"Paris wakes to life,\" encapsulates the renewed energy and vibrancy of the city. The Haiku adheres to the 5-7-5 syllable structure and portrays a vivid seasonal transformation in a concise and poetic manner. Excellent work!\n", + "\n", + "APPROVE\n", + "[Prompt tokens: 746, Completion tokens: 93]\n", + "---------- Summary ----------\n", + "Number of messages: 3\n", + "Finish reason: Text 'APPROVE' mentioned\n", + "Total prompt tokens: 1213\n", + "Total completion tokens: 112\n", + "Duration: 2.75 seconds\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Write a unique, Haiku about the weather in Paris'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=467, completion_tokens=19), content='Spring breeze gently hums, \\nCherry blossoms in full bloom— \\nParis wakes to life.'), TextMessage(source='critic', models_usage=RequestUsage(prompt_tokens=746, completion_tokens=93), content='The Haiku beautifully captures the awakening of Paris in the spring. The imagery of a gentle spring breeze and cherry blossoms in full bloom effectively conveys the rejuvenating feel of the season. The final line, \"Paris wakes to life,\" encapsulates the renewed energy and vibrancy of the city. The Haiku adheres to the 5-7-5 syllable structure and portrays a vivid seasonal transformation in a concise and poetic manner. Excellent work!\\n\\nAPPROVE')], stop_reason=\"Text 'APPROVE' mentioned\")" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "max_msg_termination = MaxMessageTermination(max_messages=10)\n", + "text_termination = TextMentionTermination(\"APPROVE\")\n", + "combined_termination = max_msg_termination | text_termination\n", + "\n", + "round_robin_team = RoundRobinGroupChat([primary_agent, critic_agent], termination_condition=combined_termination)\n", + "\n", + "# Use asyncio.run(...) if you are running this script as a standalone script.\n", + "await Console(round_robin_team.run_stream(task=\"Write a unique, Haiku about the weather in Paris\"))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The conversation stopped after the critic agent approved the message, although it could have also stopped if 10 messages were generated.\n", + "\n", + "Alternatively, if we want to stop the run only when both conditions are met, we can use the AND (`&`) operator." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "combined_termination = max_msg_termination & text_termination" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Custom Termination Condition\n", + "\n", + "The built-in termination conditions are sufficient for most use cases.\n", + "However, there may be cases where you need to implement a custom termination condition that doesn't fit into the existing ones.\n", + "You can do this by subclassing the {py:class}`~autogen_agentchat.base.TerminationCondition` class.\n", + "\n", + "In this example, we create a custom termination condition that stops the conversation when\n", + "a specific function call is made." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Sequence\n", + "\n", + "from autogen_agentchat.base import TerminatedException, TerminationCondition\n", + "from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage, StopMessage, ToolCallExecutionEvent\n", + "from autogen_core import Component\n", + "from pydantic import BaseModel\n", + "from typing_extensions import Self\n", + "\n", + "\n", + "class FunctionCallTerminationConfig(BaseModel):\n", + " \"\"\"Configuration for the termination condition to allow for serialization\n", + " and deserialization of the component.\n", + " \"\"\"\n", + "\n", + " function_name: str\n", + "\n", + "\n", + "class FunctionCallTermination(TerminationCondition, Component[FunctionCallTerminationConfig]):\n", + " \"\"\"Terminate the conversation if a FunctionExecutionResult with a specific name is received.\"\"\"\n", + "\n", + " component_config_schema = FunctionCallTerminationConfig\n", + " component_provider_override = \"autogen_agentchat.conditions.FunctionCallTermination\"\n", + " \"\"\"The schema for the component configuration.\"\"\"\n", + "\n", + " def __init__(self, function_name: str) -> None:\n", + " self._terminated = False\n", + " self._function_name = function_name\n", + "\n", + " @property\n", + " def terminated(self) -> bool:\n", + " return self._terminated\n", + "\n", + " async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:\n", + " if self._terminated:\n", + " raise TerminatedException(\"Termination condition has already been reached\")\n", + " for message in messages:\n", + " if isinstance(message, ToolCallExecutionEvent):\n", + " for execution in message.content:\n", + " if execution.name == self._function_name:\n", + " self._terminated = True\n", + " return StopMessage(\n", + " content=f\"Function '{self._function_name}' was executed.\",\n", + " source=\"FunctionCallTermination\",\n", + " )\n", + " return None\n", + "\n", + " async def reset(self) -> None:\n", + " self._terminated = False\n", + "\n", + " def _to_config(self) -> FunctionCallTerminationConfig:\n", + " return FunctionCallTerminationConfig(\n", + " function_name=self._function_name,\n", + " )\n", + "\n", + " @classmethod\n", + " def _from_config(cls, config: FunctionCallTerminationConfig) -> Self:\n", + " return cls(\n", + " function_name=config.function_name,\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Let's use this new termination condition to stop the conversation when the critic agent approves a message\n", + "using the `approve` function call.\n", + "\n", + "First we create a simple function that will be called when the critic agent approves a message." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "def approve() -> None:\n", + " \"\"\"Approve the message when all feedbacks have been addressed.\"\"\"\n", + " pass" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Then we create the agents. The critic agent is equipped with the `approve` tool." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.teams import RoundRobinGroupChat\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "model_client = OpenAIChatCompletionClient(\n", + " model=\"gpt-4o\",\n", + " temperature=1,\n", + " # api_key=\"sk-...\", # Optional if you have an OPENAI_API_KEY env variable set.\n", + ")\n", + "\n", + "# Create the primary agent.\n", + "primary_agent = AssistantAgent(\n", + " \"primary\",\n", + " model_client=model_client,\n", + " system_message=\"You are a helpful AI assistant.\",\n", + ")\n", + "\n", + "# Create the critic agent with the approve function as a tool.\n", + "critic_agent = AssistantAgent(\n", + " \"critic\",\n", + " model_client=model_client,\n", + " tools=[approve], # Register the approve function as a tool.\n", + " system_message=\"Provide constructive feedback. Use the approve tool to approve when all feedbacks are addressed.\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now, we create the termination condition and the team.\n", + "We run the team with the poem-writing task." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "Write a unique, Haiku about the weather in Paris\n", + "---------- primary ----------\n", + "Raindrops gently fall, \n", + "Cobblestones shine in dim light— \n", + "Paris dreams in grey. \n", + "---------- critic ----------\n", + "This Haiku beautifully captures a melancholic yet romantic image of Paris in the rain. The use of sensory imagery like \"Raindrops gently fall\" and \"Cobblestones shine\" effectively paints a vivid picture. It could be interesting to experiment with more distinct seasonal elements of Paris, such as incorporating the Seine River or iconic landmarks in the context of the weather. Overall, it successfully conveys the atmosphere of Paris in subtle, poetic imagery.\n", + "---------- primary ----------\n", + "Thank you for your feedback! I’m glad you enjoyed the imagery. Here’s another Haiku that incorporates iconic Parisian elements:\n", + "\n", + "Eiffel stands in mist, \n", + "Seine's ripple mirrors the sky— \n", + "Spring whispers anew. \n", + "---------- critic ----------\n", + "[FunctionCall(id='call_QEWJZ873EG4UIEpsQHi1HsAu', arguments='{}', name='approve')]\n", + "---------- critic ----------\n", + "[FunctionExecutionResult(content='None', name='approve', call_id='call_QEWJZ873EG4UIEpsQHi1HsAu', is_error=False)]\n", + "---------- critic ----------\n", + "None\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Write a unique, Haiku about the weather in Paris', type='TextMessage'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=30, completion_tokens=23), metadata={}, content='Raindrops gently fall, \\nCobblestones shine in dim light— \\nParis dreams in grey. ', type='TextMessage'), TextMessage(source='critic', models_usage=RequestUsage(prompt_tokens=99, completion_tokens=90), metadata={}, content='This Haiku beautifully captures a melancholic yet romantic image of Paris in the rain. The use of sensory imagery like \"Raindrops gently fall\" and \"Cobblestones shine\" effectively paints a vivid picture. It could be interesting to experiment with more distinct seasonal elements of Paris, such as incorporating the Seine River or iconic landmarks in the context of the weather. Overall, it successfully conveys the atmosphere of Paris in subtle, poetic imagery.', type='TextMessage'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=152, completion_tokens=48), metadata={}, content=\"Thank you for your feedback! I’m glad you enjoyed the imagery. Here’s another Haiku that incorporates iconic Parisian elements:\\n\\nEiffel stands in mist, \\nSeine's ripple mirrors the sky— \\nSpring whispers anew. \", type='TextMessage'), ToolCallRequestEvent(source='critic', models_usage=RequestUsage(prompt_tokens=246, completion_tokens=11), metadata={}, content=[FunctionCall(id='call_QEWJZ873EG4UIEpsQHi1HsAu', arguments='{}', name='approve')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='critic', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='None', name='approve', call_id='call_QEWJZ873EG4UIEpsQHi1HsAu', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='critic', models_usage=None, metadata={}, content='None', type='ToolCallSummaryMessage')], stop_reason=\"Function 'approve' was executed.\")" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "function_call_termination = FunctionCallTermination(function_name=\"approve\")\n", + "round_robin_team = RoundRobinGroupChat([primary_agent, critic_agent], termination_condition=function_call_termination)\n", + "\n", + "# Use asyncio.run(...) if you are running this script as a standalone script.\n", + "await Console(round_robin_team.run_stream(task=\"Write a unique, Haiku about the weather in Paris\"))\n", + "await model_client.close()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can see that the conversation stopped when the critic agent approved the message using the `approve` function call." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/experimental.md b/python/docs/src/user-guide/autogenstudio-user-guide/experimental.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/experimental.md rename to python/docs/src/user-guide/autogenstudio-user-guide/experimental.md diff --git a/python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/faq.md b/python/docs/src/user-guide/autogenstudio-user-guide/faq.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/faq.md rename to python/docs/src/user-guide/autogenstudio-user-guide/faq.md diff --git a/python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/index.md b/python/docs/src/user-guide/autogenstudio-user-guide/index.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/index.md rename to python/docs/src/user-guide/autogenstudio-user-guide/index.md diff --git a/python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/installation.md b/python/docs/src/user-guide/autogenstudio-user-guide/installation.md similarity index 92% rename from python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/installation.md rename to python/docs/src/user-guide/autogenstudio-user-guide/installation.md index 374ce2ffd8be..5b370d4cc045 100644 --- a/python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/installation.md +++ b/python/docs/src/user-guide/autogenstudio-user-guide/installation.md @@ -19,11 +19,18 @@ We recommend using a virtual environment as this will ensure that the dependenci Create and activate: +Linux/Mac: ```bash python3 -m venv .venv source .venv/bin/activate ``` +Windows command-line: +```batch +python3 -m venv .venv +.venv\Scripts\activate.bat +``` + To deactivate later, run: ```bash @@ -74,8 +81,9 @@ You have two options for installing from source: manually or using a dev contain ### A) Install from source manually 1. Ensure you have Python 3.10+ and Node.js (version above 14.15.0) installed. -2. Clone the AutoGen Studio repository and install its Python dependencies using `pip install -e .` -3. Navigate to the `python/packages/autogen-studio/frontend` directory, install the dependencies, and build the UI: +2. Clone the AutoGen Studio repository. +3. Navigate to the `python/packages/autogen-studio` and install its Python dependencies using `pip install -e .` +4. Navigate to the `python/packages/autogen-studio/frontend` directory, install the dependencies, and build the UI: ```bash npm install -g gatsby-cli diff --git a/python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/jsoneditor.jpg b/python/docs/src/user-guide/autogenstudio-user-guide/jsoneditor.jpg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/jsoneditor.jpg rename to python/docs/src/user-guide/autogenstudio-user-guide/jsoneditor.jpg diff --git a/python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/teambuilder.jpg b/python/docs/src/user-guide/autogenstudio-user-guide/teambuilder.jpg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/teambuilder.jpg rename to python/docs/src/user-guide/autogenstudio-user-guide/teambuilder.jpg diff --git a/python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/usage.md b/python/docs/src/user-guide/autogenstudio-user-guide/usage.md similarity index 99% rename from python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/usage.md rename to python/docs/src/user-guide/autogenstudio-user-guide/usage.md index 4d1d493630ba..fcb3065f0fcf 100644 --- a/python/packages/autogen-core/docs/src/user-guide/autogenstudio-user-guide/usage.md +++ b/python/docs/src/user-guide/autogenstudio-user-guide/usage.md @@ -50,7 +50,7 @@ AGS also lets you directly modify the JSON configuration of the team. This can b > Did you know that you define your agents in Python, export them to JSON and then paste them in the JSON editor? The section below shows how to accomplish this. -## Declarative Specification of Componenents +## Declarative Specification of Components AutoGen Studio is built on the declarative specification behaviors of AutoGen AgentChat. This allows users to define teams, agents, models, tools, and termination conditions in Python and then dump them into a JSON file for use in AutoGen Studio. diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/components/command-line-code-executors.ipynb b/python/docs/src/user-guide/core-user-guide/components/command-line-code-executors.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/components/command-line-code-executors.ipynb rename to python/docs/src/user-guide/core-user-guide/components/command-line-code-executors.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/components/model-clients.ipynb b/python/docs/src/user-guide/core-user-guide/components/model-clients.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/components/model-clients.ipynb rename to python/docs/src/user-guide/core-user-guide/components/model-clients.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/components/model-context.ipynb b/python/docs/src/user-guide/core-user-guide/components/model-context.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/components/model-context.ipynb rename to python/docs/src/user-guide/core-user-guide/components/model-context.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/components/tools.ipynb b/python/docs/src/user-guide/core-user-guide/components/tools.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/components/tools.ipynb rename to python/docs/src/user-guide/core-user-guide/components/tools.ipynb diff --git a/python/docs/src/user-guide/core-user-guide/components/workbench.ipynb b/python/docs/src/user-guide/core-user-guide/components/workbench.ipynb new file mode 100644 index 000000000000..60d527b2fa74 --- /dev/null +++ b/python/docs/src/user-guide/core-user-guide/components/workbench.ipynb @@ -0,0 +1,327 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "b2a89a30", + "metadata": {}, + "source": [ + "# Workbench (and MCP)\n", + "\n", + "A {py:class}`~autogen_core.tools.Workbench` provides a collection of tools that share state and resources.\n", + "Different from {py:class}`~autogen_core.tools.Tool`, which provides an interface\n", + "to a single tool, a workbench provides an interface to call different tools\n", + "and receive results as the same types." + ] + }, + { + "cell_type": "markdown", + "id": "f6aa6692", + "metadata": {}, + "source": [ + "## Using Workbench\n", + "\n", + "Here is an example of how to create an agent using {py:class}`~autogen_core.tools.Workbench`." + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "e8a489ec", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from dataclasses import dataclass\n", + "from typing import List\n", + "\n", + "from autogen_core import (\n", + " FunctionCall,\n", + " MessageContext,\n", + " RoutedAgent,\n", + " message_handler,\n", + ")\n", + "from autogen_core.model_context import ChatCompletionContext\n", + "from autogen_core.models import (\n", + " AssistantMessage,\n", + " ChatCompletionClient,\n", + " FunctionExecutionResult,\n", + " FunctionExecutionResultMessage,\n", + " LLMMessage,\n", + " SystemMessage,\n", + " UserMessage,\n", + ")\n", + "from autogen_core.tools import ToolResult, Workbench" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "66674f0d", + "metadata": {}, + "outputs": [], + "source": [ + "@dataclass\n", + "class Message:\n", + " content: str\n", + "\n", + "\n", + "class WorkbenchAgent(RoutedAgent):\n", + " def __init__(\n", + " self, model_client: ChatCompletionClient, model_context: ChatCompletionContext, workbench: Workbench\n", + " ) -> None:\n", + " super().__init__(\"An agent with a workbench\")\n", + " self._system_messages: List[LLMMessage] = [SystemMessage(content=\"You are a helpful AI assistant.\")]\n", + " self._model_client = model_client\n", + " self._model_context = model_context\n", + " self._workbench = workbench\n", + "\n", + " @message_handler\n", + " async def handle_user_message(self, message: Message, ctx: MessageContext) -> Message:\n", + " # Add the user message to the model context.\n", + " await self._model_context.add_message(UserMessage(content=message.content, source=\"user\"))\n", + " print(\"---------User Message-----------\")\n", + " print(message.content)\n", + "\n", + " # Run the chat completion with the tools.\n", + " create_result = await self._model_client.create(\n", + " messages=self._system_messages + (await self._model_context.get_messages()),\n", + " tools=(await self._workbench.list_tools()),\n", + " cancellation_token=ctx.cancellation_token,\n", + " )\n", + "\n", + " # Run tool call loop.\n", + " while isinstance(create_result.content, list) and all(\n", + " isinstance(call, FunctionCall) for call in create_result.content\n", + " ):\n", + " print(\"---------Function Calls-----------\")\n", + " for call in create_result.content:\n", + " print(call)\n", + "\n", + " # Add the function calls to the model context.\n", + " await self._model_context.add_message(AssistantMessage(content=create_result.content, source=\"assistant\"))\n", + "\n", + " # Call the tools using the workbench.\n", + " print(\"---------Function Call Results-----------\")\n", + " results: List[ToolResult] = []\n", + " for call in create_result.content:\n", + " result = await self._workbench.call_tool(\n", + " call.name, arguments=json.loads(call.arguments), cancellation_token=ctx.cancellation_token\n", + " )\n", + " results.append(result)\n", + " print(result)\n", + "\n", + " # Add the function execution results to the model context.\n", + " await self._model_context.add_message(\n", + " FunctionExecutionResultMessage(\n", + " content=[\n", + " FunctionExecutionResult(\n", + " call_id=call.id,\n", + " content=result.to_text(),\n", + " is_error=result.is_error,\n", + " name=result.name,\n", + " )\n", + " for call, result in zip(create_result.content, results, strict=False)\n", + " ]\n", + " )\n", + " )\n", + "\n", + " # Run the chat completion again to reflect on the history and function execution results.\n", + " create_result = await self._model_client.create(\n", + " messages=self._system_messages + (await self._model_context.get_messages()),\n", + " tools=(await self._workbench.list_tools()),\n", + " cancellation_token=ctx.cancellation_token,\n", + " )\n", + "\n", + " # Now we have a single message as the result.\n", + " assert isinstance(create_result.content, str)\n", + "\n", + " print(\"---------Final Response-----------\")\n", + " print(create_result.content)\n", + "\n", + " # Add the assistant message to the model context.\n", + " await self._model_context.add_message(AssistantMessage(content=create_result.content, source=\"assistant\"))\n", + "\n", + " # Return the result as a message.\n", + " return Message(content=create_result.content)" + ] + }, + { + "cell_type": "markdown", + "id": "1361cce4", + "metadata": {}, + "source": [ + "In this example, the agent calls the tools provided by the workbench\n", + "in a loop until the model returns a final answer." + ] + }, + { + "cell_type": "markdown", + "id": "c7f78834", + "metadata": {}, + "source": [ + "## MCP Workbench\n", + "\n", + "[Model Context Protocol (MCP)](https://modelcontextprotocol.io/) is a protocol\n", + "for providing tools and resources\n", + "to language models. An MCP server hosts a set of tools and manages their state,\n", + "while an MCP client operates from the side of the language model and\n", + "communicates with the server to access the tools, and to provide the\n", + "language model with the context it needs to use the tools effectively.\n", + "\n", + "In AutoGen, we provide {py:class}`~autogen_ext.tools.mcp.McpWorkbench`\n", + "that implements an MCP client. You can use it to create an agent that\n", + "uses tools provided by MCP servers." + ] + }, + { + "cell_type": "markdown", + "id": "ff304136", + "metadata": {}, + "source": [ + "## Web Browsing Agent using Playwright MCP\n", + "\n", + "Here is an example of how we can use the [Playwright MCP server](https://github.com/microsoft/playwright-mcp)\n", + "and the `WorkbenchAgent` class to create a web browsing agent.\n", + "\n", + "You may need to install the browser dependencies for Playwright." + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "8500959b", + "metadata": { + "vscode": { + "languageId": "shellscript" + } + }, + "outputs": [], + "source": [ + "# npx playwright install chrome" + ] + }, + { + "cell_type": "markdown", + "id": "103fa5f2", + "metadata": {}, + "source": [ + "Start the Playwright MCP server in a terminal." + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "6d1250cc", + "metadata": { + "vscode": { + "languageId": "shellscript" + } + }, + "outputs": [], + "source": [ + "# npx @playwright/mcp@latest --port 8931" + ] + }, + { + "cell_type": "markdown", + "id": "da1dcb26", + "metadata": {}, + "source": [ + "Then, create the agent using the `WorkbenchAgent` class and\n", + "{py:class}`~autogen_ext.tools.mcp.McpWorkbench` with the Playwright MCP server URL." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "578420c4", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------User Message-----------\n", + "Use Bing to find out the address of Microsoft Building 99\n", + "---------Function Calls-----------\n", + "FunctionCall(id='call_oJl0E0hWvmKZrzAM7huiIyus', arguments='{\"url\": \"https://www.bing.com\"}', name='browser_navigate')\n", + "FunctionCall(id='call_Qfab5bAsveZIVg2v0aHl4Kgv', arguments='{}', name='browser_snapshot')\n", + "---------Function Call Results-----------\n", + "type='ToolResult' name='browser_navigate' result=[TextResultContent(type='TextResultContent', content='- Ran Playwright code:\\n```js\\n// Navigate to https://www.bing.com\\nawait page.goto(\\'https://www.bing.com\\');\\n```\\n\\n- Page URL: https://www.bing.com/\\n- Page Title: Search - Microsoft Bing\\n- Page Snapshot\\n```yaml\\n- generic [ref=s1e2]:\\n - generic [ref=s1e4]:\\n - generic:\\n - generic [ref=s1e6]:\\n - generic [ref=s1e7]\\n - generic [ref=s1e10]:\\n - img \"Background image\" [ref=s1e12]\\n - generic [ref=s1e14]:\\n - generic [ref=s1e17]\\n - generic [ref=s1e18]:\\n - img \"Background image\" [ref=s1e20]\\n - main [ref=s1e23]:\\n - generic [ref=s1e24]:\\n - generic [ref=s1e25]:\\n - heading \"Trending Now on Bing\" [level=1] [ref=s1e26]\\n - navigation [ref=s1e27]:\\n - menubar [ref=s1e28]:\\n - menuitem \"Copilot\" [ref=s1e29]:\\n - link \"Copilot\" [ref=s1e30]:\\n - /url: /chat?FORM=hpcodx\\n - text: Copilot\\n - menuitem \"Images\" [ref=s1e34]:\\n - link \"Images\" [ref=s1e35]:\\n - /url: /images?FORM=Z9LH\\n - menuitem \"Videos\" [ref=s1e36]:\\n - link \"Videos\" [ref=s1e37]:\\n - /url: /videos?FORM=Z9LH1\\n - menuitem \"Shopping\" [ref=s1e38]:\\n - link \"Shopping\" [ref=s1e39]:\\n - /url: /shop?FORM=Z9LHS4\\n - menuitem \"Maps\" [ref=s1e40]:\\n - link \"Maps\" [ref=s1e41]:\\n - /url: /maps?FORM=Z9LH2\\n - menuitem \"News\" [ref=s1e42]:\\n - link \"News\" [ref=s1e43]:\\n - /url: /news/search?q=Top+stories&nvaug=%5bNewsVertical+Category%3d%22rt_MaxClass%22%5d&FORM=Z9LH3\\n - menuitem \". . . More\" [ref=s1e44]:\\n - text: . . .\\n - tooltip \"More\" [ref=s1e45]\\n - generic\\n - generic [ref=s1e49]:\\n - search [ref=s1e50]:\\n - generic [ref=s1e52]:\\n - textbox \"0 characters out of 2000\" [ref=s1e53]\\n - button \"Search using voice\" [ref=s1e55]:\\n - img [ref=s1e56]\\n - text: Search using voice\\n - link \"Open Copilot\" [ref=s1e61]:\\n - /url: /chat?FORM=hpcodx\\n - generic [ref=s1e63]\\n - generic\\n - generic [ref=s1e67]:\\n - generic [ref=s1e69]:\\n - generic [ref=s1e71]:\\n - generic:\\n - link \"Get the new Bing Wallpaper app\":\\n - /url: https://go.microsoft.com/fwlink/?linkid=2127455\\n - text: Get the new Bing Wallpaper app\\n - \\'heading \"Image of the day: Spire Cove in Kenai Fjords National Park, Seward, Alaska\" [level=3]\\':\\n - \\'link \"Image of the day: Spire Cove in Kenai Fjords National Park, Seward, Alaska\"\\':\\n - /url: /search?q=Kenai+Fjords+National+Park+Alaska&form=hpcapt&filters=HpDate:\"20250424_0700\"\\n - text: Spire Cove in Kenai Fjords National Park, Seward, Alaska\\n - generic:\\n - text: © Wander Photography/Getty Images\\n - list:\\n - listitem:\\n - button \"Download this image. Use of this image is restricted\\n to wallpaper only.\"\\n - generic [ref=s1e84]:\\n - link \"Rugged peaks and wild waters\" [ref=s1e86]:\\n - /url: /search?q=Kenai+Fjords+National+Park+Alaska&form=hpcapt&filters=HpDate:\"20250424_0700\"\\n - heading \"Rugged peaks and wild waters\" [level=2] [ref=s1e88]\\n - generic [ref=s1e89]:\\n - button \"Previous image\" [disabled] [ref=s1e90]\\n - button \"Next image\" [disabled] [ref=s1e91]\\n - button \"Feedback\" [ref=s1e92]:\\n - img [ref=s1e93]\\n - text: Feedback\\n - complementary\\n```')] is_error=False\n", + "type='ToolResult' name='browser_snapshot' result=[TextResultContent(type='TextResultContent', content='- Ran Playwright code:\\n```js\\n// \\n```\\n\\n- Page URL: https://www.bing.com/\\n- Page Title: Search - Microsoft Bing\\n- Page Snapshot\\n```yaml\\n- generic [ref=s2e2]:\\n - generic [ref=s2e4]:\\n - generic:\\n - generic [ref=s2e6]:\\n - generic [ref=s2e7]\\n - generic [ref=s2e10]:\\n - img \"Background image\" [ref=s2e12]\\n - generic [ref=s2e14]:\\n - generic [ref=s2e17]\\n - generic [ref=s2e18]:\\n - img \"Background image\" [ref=s2e20]\\n - main [ref=s2e23]:\\n - generic [ref=s2e24]:\\n - generic [ref=s2e25]:\\n - heading \"Trending Now on Bing\" [level=1] [ref=s2e26]\\n - navigation [ref=s2e27]:\\n - menubar [ref=s2e28]:\\n - menuitem \"Copilot\" [ref=s2e29]:\\n - link \"Copilot\" [ref=s2e30]:\\n - /url: /chat?FORM=hpcodx\\n - text: Copilot\\n - menuitem \"Images\" [ref=s2e34]:\\n - link \"Images\" [ref=s2e35]:\\n - /url: /images?FORM=Z9LH\\n - menuitem \"Videos\" [ref=s2e36]:\\n - link \"Videos\" [ref=s2e37]:\\n - /url: /videos?FORM=Z9LH1\\n - menuitem \"Shopping\" [ref=s2e38]:\\n - link \"Shopping\" [ref=s2e39]:\\n - /url: /shop?FORM=Z9LHS4\\n - menuitem \"Maps\" [ref=s2e40]:\\n - link \"Maps\" [ref=s2e41]:\\n - /url: /maps?FORM=Z9LH2\\n - menuitem \"News\" [ref=s2e42]:\\n - link \"News\" [ref=s2e43]:\\n - /url: /news/search?q=Top+stories&nvaug=%5bNewsVertical+Category%3d%22rt_MaxClass%22%5d&FORM=Z9LH3\\n - menuitem \". . . More\" [ref=s2e44]:\\n - text: . . .\\n - tooltip \"More\" [ref=s2e45]\\n - generic\\n - generic [ref=s2e49]:\\n - search [ref=s2e50]:\\n - generic [ref=s2e52]:\\n - textbox \"0 characters out of 2000\" [ref=s2e53]\\n - button \"Search using voice\" [ref=s2e55]:\\n - img [ref=s2e56]\\n - text: Search using voice\\n - link \"Open Copilot\" [ref=s2e61]:\\n - /url: /chat?FORM=hpcodx\\n - generic [ref=s2e63]\\n - generic\\n - generic [ref=s2e67]:\\n - generic [ref=s2e69]:\\n - generic [ref=s2e71]:\\n - generic:\\n - link \"Get the new Bing Wallpaper app\":\\n - /url: https://go.microsoft.com/fwlink/?linkid=2127455\\n - text: Get the new Bing Wallpaper app\\n - \\'heading \"Image of the day: Spire Cove in Kenai Fjords National Park, Seward, Alaska\" [level=3]\\':\\n - \\'link \"Image of the day: Spire Cove in Kenai Fjords National Park, Seward, Alaska\"\\':\\n - /url: /search?q=Kenai+Fjords+National+Park+Alaska&form=hpcapt&filters=HpDate:\"20250424_0700\"\\n - text: Spire Cove in Kenai Fjords National Park, Seward, Alaska\\n - generic:\\n - text: © Wander Photography/Getty Images\\n - list:\\n - listitem:\\n - button \"Download this image. Use of this image is restricted\\n to wallpaper only.\"\\n - generic [ref=s2e84]:\\n - link \"Rugged peaks and wild waters\" [ref=s2e86]:\\n - /url: /search?q=Kenai+Fjords+National+Park+Alaska&form=hpcapt&filters=HpDate:\"20250424_0700\"\\n - heading \"Rugged peaks and wild waters\" [level=2] [ref=s2e88]\\n - generic [ref=s2e89]:\\n - button \"Previous image\" [disabled] [ref=s2e90]\\n - button \"Next image\" [disabled] [ref=s2e91]\\n - button \"Feedback\" [ref=s2e92]:\\n - img [ref=s2e93]\\n - text: Feedback\\n - complementary\\n```')] is_error=False\n", + "---------Function Calls-----------\n", + "FunctionCall(id='call_D1X5emmqqTxiaRtCsZiGHuBr', arguments='{\"url\":\"https://www.microsoft.com\"}', name='browser_navigate')\n", + "---------Function Call Results-----------\n", + "type='ToolResult' name='browser_navigate' result=[TextResultContent(type='TextResultContent', content='- Ran Playwright code:\\n```js\\n// Navigate to https://www.microsoft.com\\nawait page.goto(\\'https://www.microsoft.com\\');\\n```\\n\\n- Page URL: https://www.microsoft.com/en-us/\\n- Page Title: Microsoft – AI, Cloud, Productivity, Computing, Gaming & Apps\\n- Page Snapshot\\n```yaml\\n- generic [ref=s1e2]:\\n - generic [ref=s1e5]:\\n - generic [ref=s1e7]:\\n - generic [ref=s1e8]:\\n - generic\\n - link \"Skip to main content\" [ref=s1e12]:\\n - /url: javascript:void(0)\\n - banner [ref=s1e13]:\\n - generic [ref=s1e15]:\\n - link \"Microsoft\" [ref=s1e16]:\\n - /url: https://www.microsoft.com\\n - navigation \"Contextual menu\" [ref=s1e17]:\\n - list [ref=s1e18]:\\n - listitem [ref=s1e19]:\\n - link \"Microsoft 365\" [ref=s1e20]:\\n - /url: https://www.microsoft.com/microsoft-365\\n - listitem [ref=s1e21]:\\n - link \"Teams\" [ref=s1e22]:\\n - /url: https://www.microsoft.com/en-us/microsoft-teams/group-chat-software\\n - listitem [ref=s1e23]:\\n - link \"Copilot\" [ref=s1e24]:\\n - /url: https://copilot.microsoft.com/\\n - listitem [ref=s1e25]:\\n - link \"Windows\" [ref=s1e26]:\\n - /url: https://www.microsoft.com/en-us/windows/\\n - listitem [ref=s1e27]:\\n - link \"Surface\" [ref=s1e28]:\\n - /url: https://www.microsoft.com/surface\\n - listitem [ref=s1e29]:\\n - link \"Xbox\" [ref=s1e30]:\\n - /url: https://www.xbox.com/\\n - listitem [ref=s1e31]:\\n - link \"Deals\" [ref=s1e32]:\\n - /url: https://www.microsoft.com/en-us/store/b/sale?icid=gm_nav_L0_salepage\\n - listitem [ref=s1e33]:\\n - link \"Small Business\" [ref=s1e34]:\\n - /url: https://www.microsoft.com/en-us/store/b/business\\n - listitem [ref=s1e35]:\\n - link \"Support\" [ref=s1e36]:\\n - /url: https://support.microsoft.com/en-us\\n - generic [ref=s1e37]:\\n - navigation \"All Microsoft menu\" [ref=s1e39]:\\n - list [ref=s1e40]:\\n - listitem [ref=s1e41]:\\n - button \"All Microsoft \\ue70d\" [ref=s1e43]:\\n - text: All Microsoft\\n - text: \\ue70d\\n - search [ref=s1e45]:\\n - button \"Search Microsoft.com\" [ref=s1e47]: \\ue721\\n - link \"0 items in shopping cart\" [ref=s1e48]:\\n - /url: https://www.microsoft.com/en-us/store/cart\\n - text: \\ue7bf\\n - generic [ref=s1e50]:\\n - link \"Sign in to your account\" [ref=s1e51]:\\n - /url: https://www.microsoft.com/cascadeauth/store/account/signin?ru=https%3A%2F%2Fwww.microsoft.com%2Fen-us%2F\\n - text: Sign in to your account\\n - generic [ref=s1e56]:\\n - main [ref=s1e58]:\\n - generic [ref=s1e59]:\\n - generic [ref=s1e62]:\\n - generic [ref=s1e64]:\\n - region \"Announcement banner\" [ref=s1e65]:\\n - paragraph [ref=s1e67]:\\n - link \"Trade in and you could get cash back. Learn more\" [ref=s1e68]:\\n - /url: https://www.microsoft.com/en-us/store/b/why-microsoft-store?icid=mscom_marcom_TS1a_WhyBuy\\n - generic [ref=s1e70]:\\n - generic [ref=s1e71]:\\n - \\'region \"featured products and announcements slideshow: navigate using the previous and next: navigate using the slide tabs\" [ref=s1e72]\\':\\n - generic [ref=s1e74]: Slide 1 of 2. Meet Surface Pro\\n - generic [ref=s1e75]:\\n - \\'link \"Skip featured products and announcements slideshow: navigate using the previous and next: navigate using the slide tabs\" [ref=s1e76]\\':\\n - /url: \"#bd5bedab-7048-4f7d-8564-09f30af30317\"\\n - generic [ref=s1e77]:\\n - generic [ref=s1e78]:\\n - button \"Pause\" [ref=s1e79]:\\n - text: Pause\\n - text: \\uf2d9\\n - button \"Previous \\ue76b\" [ref=s1e81]:\\n - text: Previous\\n - text: \\ue76b\\n - button \"Next \\ue76c\" [ref=s1e83]:\\n - text: Next\\n - text: \\ue76c\\n - region \"1 of 2\" [ref=s1e86]:\\n - generic [ref=s1e88]:\\n - generic [ref=s1e90]:\\n - img \"A Surface Pro Flex Keyboard and a Surface Pro,\\n 11th Edition, a Copilot+ PC, in the color Sapphire.\"\\n [ref=s1e95]\\n - generic [ref=s1e97]:\\n - generic [ref=s1e99]:\\n - generic [ref=s1e101]:\\n - heading \"Meet Surface Pro\" [level=1] [ref=s1e103]\\n - text: This laptop\\'s unrivalled flexibility and AI features like Live Captions\\n and Cocreator enable you to do more than you\\n ever imagined.\\n - link \"Shop Surface Pro now\" [ref=s1e106]:\\n - /url: https://www.microsoft.com/en-us/surface/devices/surface-pro-11th-edition?icid=mscom_marcom_H1a_SurfacePro11Edition_FY24SpringSurface\\n - text: Shop now\\n - text: \"End of featured products and announcements slideshow: navigate using the\\n previous and next: navigate using the slide tabs section\"\\n - generic [ref=s1e109]:\\n - generic [ref=s1e111]:\\n - generic [ref=s1e113]:\\n - navigation \"product categories\" [ref=s1e114]:\\n - list [ref=s1e115]:\\n - listitem [ref=s1e116]:\\n - link \"Shop Surface devices\" [ref=s1e118]:\\n - /url: https://www.microsoft.com/en-us/store/b/shop-all-microsoft-surface?icid=MSCOM_QL_Surface\\n - listitem [ref=s1e119]:\\n - link \"Shop Xbox games and consoles\" [ref=s1e121]:\\n - /url: https://www.microsoft.com/en-us/store/b/xbox?icid=MSCOM_QL_Xbox\\n - listitem [ref=s1e122]:\\n - link \"Shop for accessories\" [ref=s1e124]:\\n - /url: https://www.microsoft.com/en-us/store/b/accessories?icid=MSCOM_QL_Accessories\\n - listitem [ref=s1e125]:\\n - link \"Shop business products\" [ref=s1e127]:\\n - /url: https://www.microsoft.com/en-us/store/b/business?icid=MSCOM_QL_Business\\n - text: Shop for your business\\n - listitem [ref=s1e128]:\\n - link \"Find your next PC\" [ref=s1e130]:\\n - /url: https://www.microsoft.com/en-us/windows/help-me-choose?icid=MSCOM_QL_FindPC\\n - listitem [ref=s1e131]:\\n - link \"Choose your Microsoft 365\" [ref=s1e133]:\\n - /url: https://www.microsoft.com/EN-US/microsoft-365/compare-all-microsoft-365-products?icid=MSCOM_QL_M365\\n - generic [ref=s1e135]:\\n - generic [ref=s1e137]:\\n - generic [ref=s1e139]:\\n - generic [ref=s1e141]:\\n - generic [ref=s1e143]:\\n - generic [ref=s1e144]:\\n - img \"A side view of Surface Laptop for Business in the\\n color Platinum.\" [ref=s1e149]\\n - generic [ref=s1e151]: New\\n - generic [ref=s1e152]:\\n - heading \"Surface Laptop for Business, Copilot+ PC | Intel\"\\n [level=2] [ref=s1e153]\\n - text: Uncompromising power, all-day battery life,* and unique AI\\n experiences—featuring Intel® Core™ Ultra processors\\n (Series 2).\\n - generic [ref=s1e156]:\\n - link \"Shop Surface Laptop for Business.\" [ref=s1e157]:\\n - /url: https://www.microsoft.com/en-us/d/surface-laptop-for-business-copilot-pc-intel/93dzmw6q4w2b?icid=mscom_marcom_CPH1a_SurfaceLaptopForBusinessCopilotPCIntel\\n - text: Shop now\\n - generic [ref=s1e159]:\\n - generic [ref=s1e160]:\\n - img \"Red, white, blue, and black Xbox Wireless\\n Controllers\" [ref=s1e165]\\n - generic [ref=s1e167]:\\n - heading \"Xbox controllers\" [level=2] [ref=s1e168]\\n - text: Elite, wireless, adaptive—find the controller that fits your style of\\n play.\\n - generic [ref=s1e171]:\\n - link \"Shop Xbox controllers\" [ref=s1e172]:\\n - /url: https://www.microsoft.com/en-us/store/collections/XboxControllers?icid=mscom_marcom_CPH2a_XboxControllers\\n - text: Shop now\\n - generic [ref=s1e174]:\\n - generic [ref=s1e175]:\\n - img \"An Xbox Series X 2 TB Galaxy Black Special Edition, a\\n White Xbox Series X 1 TB Digital Edition and a White Xbox\\n Series S 1 TB.\" [ref=s1e180]\\n - generic [ref=s1e182]:\\n - heading \"Trade in and get up to $150 for your used\\n console\" [level=2] [ref=s1e183]\\n - text: Buy a new Xbox Series X or S and get cash back on an eligible trade-in.\\n Limited-time offer.\\n - generic [ref=s1e186]:\\n - link \"Shop Xbox consoles\" [ref=s1e187]:\\n - /url: https://www.microsoft.com/en-us/store/collections/xboxconsoles?icid=mscom_marcom_CPH3a_XboxTradeInOffer\\n - link \"Check your device\\'s eligibility\" [ref=s1e188]:\\n - /url: https://www.microsoft.com/en-us/store/b/microsoft-trade-in?icid=mscom_marcom_CPH3b_XboxTradeInOffer\\n - generic [ref=s1e190]:\\n - generic [ref=s1e191]:\\n - img \"Fresh new Xbox games featuring Dragon Ball Sparking\\n Zero, WWE2k25 and FC25.\" [ref=s1e196]\\n - generic [ref=s1e198]:\\n - heading \"Up to 70% off games\" [level=2] [ref=s1e199]\\n - text: Score spring savings on select Xbox and PC games. Sale ends April 30.\\n - generic [ref=s1e202]:\\n - link \"Shop the Xbox and PC game sale.\" [ref=s1e203]:\\n - /url: https://www.xbox.com/en-US/games/browse/spring-sale?icid=mscom_marcom_CPH4a_XboxGameSale2025\\n - text: Shop the sale\\n - generic [ref=s1e205]:\\n - generic [ref=s1e207]:\\n - generic [ref=s1e209]:\\n - generic [ref=s1e210]:\\n - generic [ref=s1e212]:\\n - img \"A Surface Pro Signature Keyboard in Sapphire with an\\n Arc Mouse in Light Grey and Slim Pen 2.\" [ref=s1e217]\\n - generic [ref=s1e219]:\\n - generic [ref=s1e221]:\\n - generic [ref=s1e223]:\\n - heading \"Made for Surface\" [level=2] [ref=s1e225]\\n - text: Find keyboards, pens, and other essentials designed to work seamlessly\\n with your Surface device.\\n - link \"Shop Surface accessories\" [ref=s1e228]:\\n - /url: https://www.microsoft.com/en-us/store/b/surface-accessories?icid=mscom_marcom_MPH1a_SurfaceAccessories\\n - generic [ref=s1e230]:\\n - generic [ref=s1e232]:\\n - generic [ref=s1e233]:\\n - heading \"For business\" [level=2] [ref=s1e235]\\n - generic [ref=s1e237]:\\n - generic [ref=s1e238]:\\n - generic [ref=s1e240]:\\n - generic [ref=s1e241]:\\n - img \"A side view of Surface Pro for Business in the\\n color Platinum.\"\\n - generic [ref=s1e248]: New\\n - generic [ref=s1e249]:\\n - heading \"Surface Pro for Business, Copilot+ PC | Intel\"\\n [level=3] [ref=s1e250]\\n - text: Ultra-versatile and built with Intel® Core™ Ultra processors (Series 2)\\n that power AI experiences to amplify your team’s\\n productivity.\\n - generic [ref=s1e253]:\\n - link \"Shop Surface Pro for Business.\" [ref=s1e254]:\\n - /url: https://www.microsoft.com/en-us/d/surface-pro-for-business-copilot-pc-intel/8qfmn9xp1rl9?icid=mscom_marcom_CPW1a_SurfaceProForBusinessCopilotPCIntel\\n - text: Shop now\\n - generic [ref=s1e256]:\\n - generic [ref=s1e257]\\n - generic [ref=s1e264]:\\n - heading \"Microsoft 365 Copilot\" [level=3] [ref=s1e265]\\n - text: Save time and focus on the things that matter most with AI in Microsoft\\n 365 for business.\\n - generic [ref=s1e268]:\\n - link \"Learn more about Microsoft 365 Copilot\" [ref=s1e269]:\\n - /url: https://www.microsoft.com/en-us/microsoft-365/copilot/business?icid=mscom_marcom_CPW2a_M365forBusiness_Copilot\\n - text: Learn more\\n - generic [ref=s1e271]:\\n - generic [ref=s1e272]:\\n - img \"A Microsoft Teams video call.\"\\n - generic [ref=s1e279]:\\n - heading \"Get Microsoft Teams for your business\"\\n [level=3] [ref=s1e280]\\n - text: Online meetings, chat, real-time collaboration, and shared cloud\\n storage—all in one place.\\n - generic [ref=s1e283]:\\n - link \"Find the right Teams plan for your business.\" [ref=s1e284]:\\n - /url: https://www.microsoft.com/en-us/microsoft-teams/small-medium-business?icid=mscom_marcom_CPW3a_TeamsForBusiness\\n - text: Find the right plan for your business\\n - generic [ref=s1e286]:\\n - generic [ref=s1e287]\\n - generic [ref=s1e294]:\\n - heading \"Join the era of AI\" [level=3] [ref=s1e295]\\n - text: Create, communicate, and code with the latest Microsoft AI solutions.\\n - generic [ref=s1e298]:\\n - link \"Explore AI solutions\" [ref=s1e299]:\\n - /url: https://www.microsoft.com/en-us/ai?icid=mscom_marcom_CPW4a_AzureAI\\n - generic [ref=s1e301]:\\n - generic [ref=s1e303]:\\n - generic [ref=s1e304]:\\n - heading \"Explore more about AI and Copilot\" [level=2]\\n [ref=s1e306]\\n - generic [ref=s1e308]:\\n - generic [ref=s1e309]:\\n - generic [ref=s1e311]:\\n - generic [ref=s1e312]:\\n - img \"collaged illustration of a woman running up an\\n escalator surrounded by stylized charts.\"\\n - generic [ref=s1e318]:\\n - heading \"How AI makes hard work easier\" [level=3]\\n [ref=s1e319]\\n - text: Dive into the surprising ways that Copilot reduces the mental effort of\\n complex tasks and enhances quality of work.\\n - generic [ref=s1e322]:\\n - link \"Uncover the details of how AI makes hard work easier.\" [ref=s1e323]:\\n - /url: https://www.microsoft.com/en-us/worklab/ai-data-drop-the-surprising-way-ai-makes-hard-work-easier?icid=mscom_marcom_CPAI1a_AIHardWorkEasier\\n - text: Uncover the details\\n - generic [ref=s1e325]:\\n - generic [ref=s1e326]:\\n - img \"Azeem Azhar.\"\\n - generic [ref=s1e332]:\\n - heading \"How AI agents are transforming work\" [level=3]\\n [ref=s1e333]\\n - text: On the WorkLab podcast, Azeem Azhar—a global thought leader—shares\\n insights on the power of deep research AI and building\\n a \"brain trust\" of agents.\\n - generic [ref=s1e336]:\\n - link \"Learn more about how AI agents are transforming work.\" [ref=s1e337]:\\n - /url: https://www.microsoft.com/en-us/worklab/podcast/azeem-azhar-on-how-ai-agents-are-transforming-work?icid=mscom_marcom_CPAI2a_WorkLabAIAgents\\n - text: Learn more\\n - generic [ref=s1e339]:\\n - generic [ref=s1e340]:\\n - img \"A multifaceted gem reflects the possibilities of\\n AI.\"\\n - generic [ref=s1e346]:\\n - heading \"Why multimodal AI matters\" [level=3]\\n [ref=s1e347]\\n - text: AI models are using images, audio, and video to solve real-world\\n challenges—like helping doctors diagnose patients or\\n meteorologists predict storms.\\n - generic [ref=s1e350]:\\n - link \"Find out more about multimodal AI.\" [ref=s1e351]:\\n - /url: https://news.microsoft.com/source/features/ai/beyond-words-ai-goes-multimodal-to-meet-you-where-you-are/?icid=mscom_marcom_CPAI3a_MultimodalAI\\n - text: Find out more\\n - generic [ref=s1e353]:\\n - generic [ref=s1e355]:\\n - generic [ref=s1e357]:\\n - \\'region \"human-interest articles and stories slideshow: navigate using the slide tabs\" [ref=s1e358]\\':\\n - generic [ref=s1e360]: Slide 1 of 2. Earth’s future in 3D\\n - generic [ref=s1e361]:\\n - \\'link \"Skip human-interest articles and stories slideshow: navigate using the slide tabs\" [ref=s1e362]\\':\\n - /url: \"#c3c99f7a-0722-484c-9b77-b90c15e84fe1\"\\n - generic [ref=s1e363]:\\n - generic [ref=s1e364]:\\n - button \"Pause\" [ref=s1e365]:\\n - text: Pause\\n - text: \\uf2d9\\n - button \"Previous \\ue76b\" [ref=s1e367]:\\n - text: Previous\\n - text: \\ue76b\\n - button \"Next \\ue76c\" [ref=s1e369]:\\n - text: Next\\n - text: \\ue76c\\n - region \"1 of 2\" [ref=s1e372]:\\n - generic [ref=s1e374]:\\n - generic [ref=s1e376]:\\n - img \"A boy wearing a Hololens, a mixed reality\\n headset, comes face to face with a sea turtle in a\\n museum hall.\" [ref=s1e381]\\n - generic [ref=s1e383]:\\n - generic [ref=s1e385]:\\n - generic [ref=s1e387]:\\n - heading \"Earth’s future in 3D\" [level=2]\\n [ref=s1e389]\\n - text: Microsoft and the Natural History Museum London are imagining what’s\\n possible for the planet in 2125 through an\\n innovative exhibit.\\n - \\'link \"Explore Visions of Nature: A Mixed Reality Experience.\" [ref=s1e392]\\':\\n - /url: https://unlocked.microsoft.com/nhm-visions-of-nature/?icid=mscom_marcom_SAM1a_NaturalHistoryMuseum\\n - text: Explore Visions of Nature\\n - text: \"End of human-interest articles and stories slideshow: navigate using the\\n slide tabs section\"\\n - generic [ref=s1e395]:\\n - generic [ref=s1e397]:\\n - generic [ref=s1e399]:\\n - region \"follow us on social media\" [ref=s1e400]:\\n - heading \"Follow Microsoft\" [level=2] [ref=s1e401]\\n - list [ref=s1e402]:\\n - listitem [ref=s1e403]:\\n - link \"Follow Microsoft on Facebook, opens in a new tab\" [ref=s1e404]:\\n - /url: https://www.facebook.com/Microsoft\\n - img \"Facebook\" [ref=s1e405]\\n - listitem [ref=s1e406]:\\n - link \"Follow Microsoft on X, opens in a new tab\" [ref=s1e407]:\\n - /url: https://twitter.com/microsoft\\n - img \"X\" [ref=s1e408]\\n - listitem [ref=s1e409]:\\n - link \"Follow Microsoft on Linkedin, opens in a new tab\" [ref=s1e410]:\\n - /url: https://www.linkedin.com/company/microsoft\\n - img \"LinkedIn\" [ref=s1e411]\\n - generic\\n - generic:\\n - generic\\n - generic [ref=s1e420]:\\n - generic [ref=s1e421]:\\n - link \"Back to top\" [ref=s1e424]:\\n - /url: \"#page-top\"\\n - generic [ref=s1e425]:\\n - text: \\ue74a\\n - text: Back to top\\n - generic [ref=s1e428]:\\n - generic [ref=s1e430]:\\n - contentinfo [ref=s1e431]:\\n - navigation \"Footer Resource links\" [ref=s1e432]:\\n - generic:\\n - generic [ref=s1e434]:\\n - heading \"What\\'s new\" [level=2] [ref=s1e435]\\n - list [ref=s1e436]:\\n - listitem [ref=s1e437]:\\n - link \"Surface Pro What\\'s new\" [ref=s1e438]:\\n - /url: https://www.microsoft.com/en-us/surface/devices/surface-pro-11th-edition\\n - text: Surface Pro\\n - listitem [ref=s1e439]:\\n - link \"Surface Laptop What\\'s new\" [ref=s1e440]:\\n - /url: https://www.microsoft.com/en-us/surface/devices/surface-laptop-7th-edition\\n - text: Surface Laptop\\n - listitem [ref=s1e441]:\\n - link \"Surface Laptop Studio 2 What\\'s new\" [ref=s1e442]:\\n - /url: https://www.microsoft.com/en-us/d/Surface-Laptop-Studio-2/8rqr54krf1dz\\n - text: Surface Laptop Studio 2\\n - listitem [ref=s1e443]:\\n - link \"Surface Laptop Go 3 What\\'s new\" [ref=s1e444]:\\n - /url: https://www.microsoft.com/en-us/d/Surface-Laptop-Go-3/8p0wwgj6c6l2\\n - text: Surface Laptop Go 3\\n - listitem [ref=s1e445]:\\n - link \"Microsoft Copilot What\\'s new\" [ref=s1e446]:\\n - /url: https://www.microsoft.com/en-us/microsoft-copilot\\n - text: Microsoft Copilot\\n - listitem [ref=s1e447]:\\n - link \"AI in Windows What\\'s new\" [ref=s1e448]:\\n - /url: https://www.microsoft.com/en-us/windows/copilot-ai-features\\n - text: AI in Windows\\n - listitem [ref=s1e449]:\\n - link \"Explore Microsoft products What\\'s new\" [ref=s1e450]:\\n - /url: https://www.microsoft.com/en-us/microsoft-products-and-apps\\n - text: Explore Microsoft products\\n - listitem [ref=s1e451]:\\n - link \"Windows 11 apps What\\'s new\" [ref=s1e452]:\\n - /url: https://www.microsoft.com/windows/windows-11-apps\\n - text: Windows 11 apps\\n - generic [ref=s1e453]:\\n - heading \"Microsoft Store\" [level=2] [ref=s1e454]\\n - list [ref=s1e455]:\\n - listitem [ref=s1e456]:\\n - link \"Account profile Microsoft Store\" [ref=s1e457]:\\n - /url: https://account.microsoft.com/\\n - text: Account profile\\n - listitem [ref=s1e458]:\\n - link \"Download Center Microsoft Store\" [ref=s1e459]:\\n - /url: https://www.microsoft.com/en-us/download\\n - text: Download Center\\n - listitem [ref=s1e460]:\\n - link \"Microsoft Store support Microsoft Store\" [ref=s1e461]:\\n - /url: https://go.microsoft.com/fwlink/?linkid=2139749\\n - text: Microsoft Store support\\n - listitem [ref=s1e462]:\\n - link \"Returns Microsoft Store\" [ref=s1e463]:\\n - /url: https://www.microsoft.com/en-us/store/b/returns\\n - text: Returns\\n - listitem [ref=s1e464]:\\n - link \"Order tracking Microsoft Store\" [ref=s1e465]:\\n - /url: https://www.microsoft.com/en-us/store/b/order-tracking\\n - text: Order tracking\\n - listitem [ref=s1e466]:\\n - link \"Certified Refurbished Microsoft Store\" [ref=s1e467]:\\n - /url: https://www.microsoft.com/en-us/store/b/certified-refurbished-products\\n - text: Certified Refurbished\\n - listitem [ref=s1e468]:\\n - link \"Microsoft Store Promise Microsoft Store\" [ref=s1e469]:\\n - /url: https://www.microsoft.com/en-us/store/b/why-microsoft-store?icid=footer_why-msft-store_7102020\\n - text: Microsoft Store Promise\\n - listitem [ref=s1e470]:\\n - link \"Flexible Payments Microsoft Store\" [ref=s1e471]:\\n - /url: https://www.microsoft.com/en-us/store/b/payment-financing-options?icid=footer_financing_vcc\\n - text: Flexible Payments\\n - generic [ref=s1e472]:\\n - heading \"Education\" [level=2] [ref=s1e473]\\n - list [ref=s1e474]:\\n - listitem [ref=s1e475]:\\n - link \"Microsoft in education Education\" [ref=s1e476]:\\n - /url: https://www.microsoft.com/en-us/education\\n - text: Microsoft in education\\n - listitem [ref=s1e477]:\\n - link \"Devices for education Education\" [ref=s1e478]:\\n - /url: https://www.microsoft.com/en-us/education/devices/overview\\n - text: Devices for education\\n - listitem [ref=s1e479]:\\n - link \"Microsoft Teams for Education Education\" [ref=s1e480]:\\n - /url: https://www.microsoft.com/en-us/education/products/teams\\n - text: Microsoft Teams for Education\\n - listitem [ref=s1e481]:\\n - link \"Microsoft 365 Education Education\" [ref=s1e482]:\\n - /url: https://www.microsoft.com/en-us/education/products/microsoft-365\\n - text: Microsoft 365 Education\\n - listitem [ref=s1e483]:\\n - link \"How to buy for your school Education\" [ref=s1e484]:\\n - /url: https://www.microsoft.com/education/how-to-buy\\n - text: How to buy for your school\\n - listitem [ref=s1e485]:\\n - link \"Educator training and development Education\" [ref=s1e486]:\\n - /url: https://education.microsoft.com/\\n - text: Educator training and development\\n - listitem [ref=s1e487]:\\n - link \"Deals for students and parents Education\" [ref=s1e488]:\\n - /url: https://www.microsoft.com/en-us/store/b/education\\n - text: Deals for students and parents\\n - listitem [ref=s1e489]:\\n - link \"Azure for students Education\" [ref=s1e490]:\\n - /url: https://azure.microsoft.com/en-us/free/students/\\n - text: Azure for students\\n - generic:\\n - generic [ref=s1e492]:\\n - heading \"Business\" [level=2] [ref=s1e493]\\n - list [ref=s1e494]:\\n - listitem [ref=s1e495]:\\n - link \"Microsoft Cloud Business\" [ref=s1e496]:\\n - /url: https://www.microsoft.com/en-us/microsoft-cloud\\n - text: Microsoft Cloud\\n - listitem [ref=s1e497]:\\n - link \"Microsoft Security Business\" [ref=s1e498]:\\n - /url: https://www.microsoft.com/en-us/security\\n - text: Microsoft Security\\n - listitem [ref=s1e499]:\\n - link \"Dynamics 365 Business\" [ref=s1e500]:\\n - /url: https://www.microsoft.com/en-us/dynamics-365\\n - text: Dynamics 365\\n - listitem [ref=s1e501]:\\n - link \"Microsoft 365 Business\" [ref=s1e502]:\\n - /url: https://www.microsoft.com/en-us/microsoft-365/business\\n - text: Microsoft 365\\n - listitem [ref=s1e503]:\\n - link \"Microsoft Power Platform Business\" [ref=s1e504]:\\n - /url: https://www.microsoft.com/en-us/power-platform\\n - text: Microsoft Power Platform\\n - listitem [ref=s1e505]:\\n - link \"Microsoft Teams Business\" [ref=s1e506]:\\n - /url: https://www.microsoft.com/en-us/microsoft-teams/group-chat-software\\n - text: Microsoft Teams\\n - listitem [ref=s1e507]:\\n - link \"Microsoft 365 Copilot Business\" [ref=s1e508]:\\n - /url: https://www.microsoft.com/en-us/microsoft-365/copilot/copilot-for-work\\n - text: Microsoft 365 Copilot\\n - listitem [ref=s1e509]:\\n - link \"Small Business Business\" [ref=s1e510]:\\n - /url: https://www.microsoft.com/en-us/store/b/business?icid=CNavBusinessStore\\n - text: Small Business\\n - generic [ref=s1e511]:\\n - heading \"Developer & IT\" [level=2] [ref=s1e512]\\n - list [ref=s1e513]:\\n - listitem [ref=s1e514]:\\n - link \"Azure Developer & IT\" [ref=s1e515]:\\n - /url: https://azure.microsoft.com/en-us/\\n - text: Azure\\n - listitem [ref=s1e516]:\\n - link \"Microsoft Developer Developer & IT\" [ref=s1e517]:\\n - /url: https://developer.microsoft.com/en-us/\\n - text: Microsoft Developer\\n - listitem [ref=s1e518]:\\n - link \"Microsoft Learn Developer & IT\" [ref=s1e519]:\\n - /url: https://learn.microsoft.com/\\n - text: Microsoft Learn\\n - listitem [ref=s1e520]:\\n - link \"Support for AI marketplace apps Developer & IT\" [ref=s1e521]:\\n - /url: https://www.microsoft.com/isv/isv-success?ocid=cmm3atxvn98\\n - text: Support for AI marketplace apps\\n - listitem [ref=s1e522]:\\n - link \"Microsoft Tech Community Developer & IT\" [ref=s1e523]:\\n - /url: https://techcommunity.microsoft.com/\\n - text: Microsoft Tech Community\\n - listitem [ref=s1e524]:\\n - link \"Azure Marketplace Developer & IT\" [ref=s1e525]:\\n - /url: https://azuremarketplace.microsoft.com/en-us/\\n - text: Azure Marketplace\\n - listitem [ref=s1e526]:\\n - link \"AppSource Developer & IT\" [ref=s1e527]:\\n - /url: https://appsource.microsoft.com/en-us/\\n - text: AppSource\\n - listitem [ref=s1e528]:\\n - link \"Visual Studio Developer & IT\" [ref=s1e529]:\\n - /url: https://visualstudio.microsoft.com/\\n - text: Visual Studio\\n - generic [ref=s1e530]:\\n - heading \"Company\" [level=2] [ref=s1e531]\\n - list [ref=s1e532]:\\n - listitem [ref=s1e533]:\\n - link \"Careers Company\" [ref=s1e534]:\\n - /url: https://careers.microsoft.com/\\n - text: Careers\\n - listitem [ref=s1e535]:\\n - link \"About Microsoft Company\" [ref=s1e536]:\\n - /url: https://www.microsoft.com/about\\n - text: About Microsoft\\n - listitem [ref=s1e537]:\\n - link \"Company news Company\" [ref=s1e538]:\\n - /url: https://news.microsoft.com/\\n - text: Company news\\n - listitem [ref=s1e539]:\\n - link \"Privacy at Microsoft Company\" [ref=s1e540]:\\n - /url: https://privacy.microsoft.com/en-us\\n - text: Privacy at Microsoft\\n - listitem [ref=s1e541]:\\n - link \"Investors Company\" [ref=s1e542]:\\n - /url: https://www.microsoft.com/investor/default.aspx\\n - text: Investors\\n - listitem [ref=s1e543]:\\n - link \"Diversity and inclusion Company\" [ref=s1e544]:\\n - /url: https://www.microsoft.com/en-us/diversity/\\n - text: Diversity and inclusion\\n - listitem [ref=s1e545]:\\n - link \"Accessibility Company\" [ref=s1e546]:\\n - /url: https://www.microsoft.com/en-us/accessibility\\n - text: Accessibility\\n - listitem [ref=s1e547]:\\n - link \"Sustainability Company\" [ref=s1e548]:\\n - /url: https://www.microsoft.com/en-us/sustainability/\\n - text: Sustainability\\n - generic [ref=s1e549]:\\n - link \"Content Language Selector. Currently set to English (United States)\" [ref=s1e550]:\\n - /url: https://www.microsoft.com/en-us/locale\\n - text: \\ue909 English (United States)\\n - link \"Your Privacy Choices Opt-Out Icon Your Privacy Choices\" [ref=s1e551]:\\n - /url: https://aka.ms/yourcaliforniaprivacychoices\\n - img \"Your Privacy Choices Opt-Out Icon\" [ref=s1e552]\\n - text: Your Privacy Choices\\n - link \"Consumer Health Privacy\" [ref=s1e558]:\\n - /url: https://go.microsoft.com/fwlink/?linkid=2259814\\n - text: Consumer Health Privacy\\n - navigation \"Microsoft corporate links\":\\n - list [ref=s1e561]:\\n - listitem [ref=s1e562]:\\n - link \"Sitemap\" [ref=s1e563]:\\n - /url: https://www.microsoft.com/en-us/sitemap1.aspx\\n - listitem [ref=s1e564]:\\n - link \"Contact Microsoft\" [ref=s1e565]:\\n - /url: https://support.microsoft.com/contactus\\n - listitem [ref=s1e566]:\\n - link \"Privacy\" [ref=s1e567]:\\n - /url: https://go.microsoft.com/fwlink/?LinkId=521839\\n - listitem [ref=s1e568]:\\n - link \"Terms of use\" [ref=s1e569]:\\n - /url: https://go.microsoft.com/fwlink/?LinkID=206977\\n - listitem [ref=s1e570]:\\n - link \"Trademarks\" [ref=s1e571]:\\n - /url: https://go.microsoft.com/fwlink/?linkid=2196228\\n - listitem [ref=s1e572]:\\n - link \"Safety & eco\" [ref=s1e573]:\\n - /url: https://go.microsoft.com/fwlink/?linkid=2196227\\n - listitem [ref=s1e574]:\\n - link \"Recycling\" [ref=s1e575]:\\n - /url: https://www.microsoft.com/en-us/legal/compliance/recycling\\n - listitem [ref=s1e576]:\\n - link \"About our ads\" [ref=s1e577]:\\n - /url: https://choice.microsoft.com\\n - listitem [ref=s1e578]: © Microsoft 2025\\n - region \"Chat with an Expert\":\\n - generic:\\n - generic [ref=s1e587]:\\n - paragraph [ref=s1e589]:\\n - text: Need help?\\n - text: Let\\'s chat\\n - img \"Need Help? Lets Chat\" [ref=s1e591]\\n - button \"Need help? Let\\'s chat\" [ref=s1e592]\\n```')] is_error=False\n", + "---------Function Calls-----------\n", + "FunctionCall(id='call_4AtqCm5GVIRUqgR8LtJ4pGWF', arguments='{\"url\":\"https://www.bing.com/search?q=Microsoft+Building+99+address\"}', name='browser_navigate')\n", + "---------Function Call Results-----------\n", + "type='ToolResult' name='browser_navigate' result=[TextResultContent(type='TextResultContent', content='- Ran Playwright code:\\n```js\\n// Navigate to https://www.bing.com/search?q=Microsoft+Building+99+address\\nawait page.goto(\\'https://www.bing.com/search?q=Microsoft+Building+99+address\\');\\n```\\n\\n- Page URL: https://www.bing.com/search?q=Microsoft+Building+99+address\\n- Page Title: Microsoft Building 99 address - Search\\n- Page Snapshot\\n```yaml\\n- generic [ref=s1e2]:\\n - banner [ref=s1e3]:\\n - button \"Skip to content\" [ref=s1e4]:\\n - generic [ref=s1e6]: Skip to content\\n - generic [ref=s1e7]:\\n - link \"Back to Bing search\" [ref=s1e8]:\\n - /url: /?FORM=Z9FD1\\n - heading \"Back to Bing search\" [level=1] [ref=s1e9]\\n - search [ref=s1e10]:\\n - link \"Search button\" [ref=s1e12]:\\n - /url: javascript:void(0)\\n - generic [ref=s1e13]:\\n - button \"Search\" [ref=s1e15]\\n - searchbox \"Enter your search here - Search suggestions will show as you type\" [ref=s1e16]: Microsoft Building 99 address\\n - generic [ref=s1e17]\\n - generic [ref=s1e21]:\\n - button \"Search using an image\" [ref=s1e22]\\n - link \"Chat with Copilot\" [ref=s1e25]:\\n - /url: /chat?q=Microsoft+Building+99+address&sendquery=1&form=HECODX\\n - button \"Chat with Copilot\" [ref=s1e26]\\n - complementary \"Account Rewards and Preferences\" [ref=s1e28]:\\n - link \"Sign in\" [ref=s1e29]:\\n - /url: javascript:void(0)\\n - generic [ref=s1e31]:\\n - button \"Sign in\" [ref=s1e32]\\n - button \"Microsoft Rewards\" [ref=s1e33]:\\n - generic [ref=s1e35]:\\n - text: Rewards\\n - img [ref=s1e38]\\n - button \"Mobile\" [ref=s1e42]:\\n - text: Mobile\\n - img [ref=s1e44]\\n - button \"Settings and quick links\" [ref=s1e46]\\n - navigation \"Search Filter\" [ref=s1e47]:\\n - list [ref=s1e48]:\\n - listitem [ref=s1e49]:\\n - link \"All\" [ref=s1e50]:\\n - /url: /?scope=web&FORM=HDRSC1\\n - text: All\\n - listitem [ref=s1e52]:\\n - link \"Search\" [ref=s1e53]:\\n - /url: /copilotsearch?q=Microsoft+Building+99+address&FORM=CSSCOP\\n - img [ref=s1e54]\\n - text: Search\\n - listitem [ref=s1e56]:\\n - link \"Copilot\" [ref=s1e57]:\\n - /url: /chat?q=Microsoft+Building+99+address&sendquery=1&FORM=SCCODX\\n - listitem [ref=s1e58]:\\n - link \"Videos\" [ref=s1e59]:\\n - /url: /videos/search?q=Microsoft+Building+99+address&FORM=HDRSC4\\n - listitem [ref=s1e60]:\\n - link \"Images\" [ref=s1e61]:\\n - /url: /images/search?q=Microsoft+Building+99+address&FORM=HDRSC3\\n - listitem [ref=s1e62]:\\n - link \"Maps\" [ref=s1e63]:\\n - /url: /maps?q=Microsoft+Building+99+address&FORM=HDRSC6\\n - listitem [ref=s1e64]:\\n - link \"News\" [ref=s1e65]:\\n - /url: /news/search?q=Microsoft+Building+99+address&FORM=HDRSC7\\n - listitem [ref=s1e66]:\\n - button \"More\" [ref=s1e67]:\\n - img [ref=s1e69]\\n - text: More\\n - listitem [ref=s1e72]:\\n - link \"Tools\" [ref=s1e73]:\\n - /url: javascript:void(0)\\n - main \"Search Results\" [ref=s1e76]:\\n - generic [ref=s1e77]:\\n - list [ref=s1e78]:\\n - listitem [ref=s1e79]:\\n - generic [ref=s1e81]:\\n - generic [ref=s1e83]:\\n - generic [ref=s1e85]:\\n - generic [ref=s1e87]:\\n - list \"Please use arrow keys to navigate\" [ref=s1e88]:\\n - listitem [ref=s1e89]:\\n - generic [ref=s1e90]:\\n - link \"campusbuilding.com\" [ref=s1e92]:\\n - /url: https://campusbuilding.com/b/microsoft-building-99/\\n - generic [ref=s1e94]:\\n - generic [ref=s1e96]\\n - generic [ref=s1e97]:\\n - text: campusbuilding.com\\n - generic [ref=s1e100]: https://campusbuilding.com\\n - heading \"Microsoft Building 99 Building Details\" [level=2] [ref=s1e102]:\\n - link \"Microsoft Building 99 Building Details\" [ref=s1e103]:\\n - /url: https://campusbuilding.com/b/microsoft-building-99/\\n - list [ref=s1e105]:\\n - listitem [ref=s1e106]:\\n - generic [ref=s1e107]:\\n - link \"The address of Microsoft Building 99 is 14820 NE 36th St, Redmond WA 98052. Microsoft Building 99 is near the intersection of Northeast 33rd Court and 143rd Place Northeast. Micr…\" [ref=s1e109]:\\n - /url: https://campusbuilding.com/b/microsoft-building-99/\\n - text: The address of Microsoft Building 99\\n - strong [ref=s1e110]: is\\n - strong [ref=s1e111]: \"14820\"\\n - strong [ref=s1e112]: NE\\n - strong [ref=s1e113]: 36th\\n - strong [ref=s1e114]: St\\n - text: \",\"\\n - strong [ref=s1e115]: Redmond\\n - strong [ref=s1e116]: WA\\n - strong [ref=s1e117]: \"98052\"\\n - text: . Microsoft Building 99 is near the intersection of Northeast 33rd Court\\n and 143rd Place Northeast. Micr…\\n - generic [ref=s1e118]:\\n - generic [ref=s1e119]:\\n - generic [ref=s1e120]:\\n - link \"Microsoft The Commons Mixer\" [ref=s1e121]:\\n - /url: https://campusbuilding.com/b/microsoft-the-commons-mixer/\\n - text: Microsoft The Commons Mixer\\n - text: This building has a Microsoft IT Tech Link. The Microsoft Techlink is a\\n place f…\\n - generic [ref=s1e124]:\\n - link \"Microsoft Studio H\" [ref=s1e125]:\\n - /url: https://campusbuilding.com/b/microsoft-studio-h/\\n - text: Microsoft Studio H\\n - text: Food, coffee, and restaurants close to Microsoft Studio H. Microsoft Cafe\\n H is …\\n - generic [ref=s1e128]:\\n - link \"Microsoft Studio G\" [ref=s1e129]:\\n - /url: https://campusbuilding.com/b/microsoft-studio-g/\\n - text: Microsoft Studio G\\n - text: Food, coffee, and restaurants close to Microsoft Studio G. Microsoft Cafe\\n H is …\\n - generic [ref=s1e132]:\\n - generic [ref=s1e133]:\\n - link \"Microsoft Studio E\" [ref=s1e134]:\\n - /url: https://campusbuilding.com/b/microsoft-studio-e/\\n - text: Microsoft Studio E\\n - text: Microsoft Building 123 0.12 miles; Microsoft The Commons Mixer 0.12 mil…\\n - generic [ref=s1e137]:\\n - link \"Microsoft Building 113\" [ref=s1e138]:\\n - /url: https://campusbuilding.com/b/microsoft-building-113/\\n - text: Microsoft Building 113\\n - text: The address of Microsoft Building 113 is 14870 NE 31st Way, Redmond WA\\n 980…\\n - generic [ref=s1e141]:\\n - link \"Redmond Main Campus\" [ref=s1e142]:\\n - /url: https://campusbuilding.com/c/microsoft-redmond-main-campus/\\n - text: Redmond Main Campus\\n - text: There are 95 buildings at the Microsoft Redmond Main Campus.\\n - listitem [ref=s1e145]:\\n - generic [ref=s1e147]:\\n - generic [ref=s1e148]:\\n - link \"Redmond, Washington - Wikipedia\" [ref=s1e149]:\\n - /url: https://en.wikipedia.org/wiki/Redmond,_Washington\\n - heading \"Redmond, Washington - Wikipedia\"\\n [level=1] [ref=s1e150]\\n - link \"Redmond, Washington - Wikipedia\" [ref=s1e151]:\\n - /url: https://en.wikipedia.org/wiki/Redmond,_Washington\\n - text: City in Washington\\n - text: Redmond is a city in King County, Washington, United States, located 15\\n miles east of Seattle. The population was 73,256\\n at the 2020 census. Redmond is best known as the\\n home of Microsoft and Nintendo of America. The\\n city has a large technology industry in addition\\n to being a...\\n - text: See more on Wikipedia\\n - link \"Redmond, Washington - Wikipedia\" [ref=s1e156]:\\n - /url: https://en.wikipedia.org/wiki/Redmond,_Washington\\n - generic [ref=s1e158]\\n - generic [ref=s1e160]:\\n - \\'link \"Microsoft\\'\\'s Building 99 from YouTube · Duration: 44 seconds · 28.7K views · uploaded on May 20, 2010 · uploaded by CNET · Click to play.\" [ref=s1e162]\\':\\n - /url: /videos/riverview/relatedvideo?q=Microsoft+Building+99+address&&mid=15C1FEC0FBDB1C2218B715C1FEC0FBDB1C2218B7&FORM=VAMGZC\\n - generic [ref=s1e163]:\\n - generic [ref=s1e164]:\\n - img \"Microsoft\\'s Building 99\" [ref=s1e166]\\n - generic [ref=s1e169]:\\n - generic [ref=s1e171]: 00:44\\n - generic [ref=s1e172]:\\n - generic [ref=s1e174]:\\n - generic [ref=s1e176]: YouTube\\n - text: › CNET\\n - text: · 28.7K views\\n - text: · May 20, 2010\\n - generic [ref=s1e181]:\\n - link \"Microsoft Studio H Building There are at least 441 amenities within 1 mile of Microsoft Studio H. Here\\'s a summary of th… campusbuilding.com\" [ref=s1e182]:\\n - /url: https://campusbuilding.com/b/microsoft-studio-h\\n - generic [ref=s1e183]:\\n - generic [ref=s1e184]:\\n - text: Microsoft Studio H Building\\n - contentinfo [ref=s1e186]: There are at least 441 amenities within 1 mile of\\n Microsoft Studio H. Here\\'s a summary of th…\\n - generic [ref=s1e187]:\\n - generic [ref=s1e189]\\n - generic [ref=s1e192]: campusbuilding.com\\n - generic [ref=s1e193]:\\n - generic [ref=s1e195]:\\n - text: Feedback\\n - button \"Feedback Like\" [ref=s1e197]\\n - button \"Feedback Dislike\" [ref=s1e198]\\n - list [ref=s1e199]:\\n - listitem [ref=s1e200]:\\n - link \"Microsoft\" [ref=s1e202]:\\n - /url: https://www.microsoft.com/en-us/research/lab/microsoft-research-redmond/\\n - generic [ref=s1e204]:\\n - generic [ref=s1e206]\\n - generic [ref=s1e207]:\\n - text: Microsoft\\n - generic [ref=s1e210]: https://www.microsoft.com › en-us › research › lab › ...\\n - heading \"Microsoft Research Lab - Redmond - Microsoft Research\" [level=2] [ref=s1e212]:\\n - link \"Microsoft Research Lab - Redmond - Microsoft Research\" [ref=s1e213]:\\n - /url: https://www.microsoft.com/en-us/research/lab/microsoft-research-redmond/\\n - generic [ref=s1e214]:\\n - generic [ref=s1e216]:\\n - list [ref=s1e218]:\\n - listitem [ref=s1e219]:\\n - generic [ref=s1e221]:\\n - link \"Web page related images\" [ref=s1e222]:\\n - /url: /images/search?view=detailV2&ccid=wloBxYbF&id=2935EA20AFFDBD8BE5408325977F59B9C223BE65&thid=OIP.wloBxYbFlyxykavHItPjrwHaEK&mediaurl=https://www.microsoft.com/en-us/research/uploads/prod/2019/09/Jina-Shuh_Podcast_Site_09_2019_1400x788-1280x720.jpg&q=Microsoft\\n Building 99\\n address&ck=2A955D85573E45B5282A0D14F641B1D3&idpp=rc&idpview=singleimage&form=rc2idp\\n - listitem [ref=s1e224]:\\n - generic [ref=s1e226]:\\n - link \"Web page related images\" [ref=s1e227]:\\n - /url: /images/search?view=detailV2&ccid=5umtXZtt&id=2935EA20AFFDBD8BE5409E4347300E10FF614100&thid=OIP.5umtXZttL9GCgDCHUvcXMAHaEK&mediaurl=https://www.microsoft.com/en-us/research/wp-content/uploads/2024/06/RF-Ep3-Recap-BlogHeroFeature-1400x788-1.jpg&q=Microsoft\\n Building 99\\n address&ck=200A413213746B198D7ECBAB99D78F6D&idpp=rc&idpview=singleimage&form=rc2idp\\n - listitem [ref=s1e229]:\\n - generic [ref=s1e231]:\\n - link \"Web page related images\" [ref=s1e232]:\\n - /url: /images/search?view=detailV2&ccid=n+L4YRU1&id=2935EA20AFFDBD8BE5405B2B8CFB4C4B3275FBF5&thid=OIP.n-L4YRU1nH-mXzimQR4yiwHaEK&mediaurl=https://www.microsoft.com/en-us/research/uploads/prod/2023/10/Podcast_Insights_Madeline_Hero_Feature_No_Text_1400x788-960x540.png&q=Microsoft\\n Building 99\\n address&ck=EAED05EF011C3035948F31A2E52BDBF1&idpp=rc&idpview=singleimage&form=rc2idp\\n - listitem [ref=s1e234]:\\n - generic [ref=s1e236]:\\n - link \"Web page related images\" [ref=s1e237]:\\n - /url: /images/search?view=detailV2&ccid=0j+0bcqx&id=2935EA20AFFDBD8BE5406AB81021186B77C5538D&thid=OIP.0j-0bcqxnBJ56WIEnV5eEgAAAA&mediaurl=https://www.microsoft.com/en-us/research/uploads/prod/2019/05/HUE_header_04_2019_1920x720-343x193.jpg&q=Microsoft\\n Building 99\\n address&ck=7DEDB051EC61BF82B729E3244983484A&idpp=rc&idpview=singleimage&form=rc2idp\\n - paragraph [ref=s1e239]:\\n - text: Mar 7, 2025\\n - text: · Corporate Vice President and Managing Director, Microsoft Research\\n Redmond Address Microsoft Building 99, 14820 NE 36th Street,\\n Redmond, Washington, 98052 USA\\n - listitem [ref=s1e241]:\\n - link \"campusbuilding.com\" [ref=s1e243]:\\n - /url: https://campusbuilding.com/c/microsoft-redmond-main-campus/\\n - generic [ref=s1e245]:\\n - generic [ref=s1e247]\\n - generic [ref=s1e248]:\\n - text: campusbuilding.com\\n - generic [ref=s1e251]: https://campusbuilding.com › microsoft-re…\\n - generic [ref=s1e254]:\\n - generic [ref=s1e256]:\\n - link \"/images/search?view=detailV2&ccid=cbwg8zEM&id=F7D283C10E5E7DCF7A1F79F6E84887B5FDD2AC1D&thid=OIP.cbwg8zEM_3qR98xadB9dSAHaHQ&mediaurl=https://campusbuilding.com/static/images/map_of_microsoft_redmond_main_campus_and_buildings.jpg&q=Microsoft+Building+99+address&ck=F76D45DFC43E694CFE9C00250746F505&idpp=rc&idpview=singleimage&form=rc2idp&mode=overlay\" [ref=s1e257]:\\n - /url: javascript:void(0)\\n - heading \"Microsoft Redmond Main Campus and Buildings\" [level=2] [ref=s1e259]:\\n - link \"Microsoft Redmond Main Campus and Buildings\" [ref=s1e260]:\\n - /url: https://campusbuilding.com/c/microsoft-redmond-main-campus/\\n - paragraph [ref=s1e261]: There are 95 buildings at the Microsoft Redmond Main Campus.\\n - listitem [ref=s1e263]:\\n - link \"Mapcarta\" [ref=s1e265]:\\n - /url: https://mapcarta.com/W93639217\\n - generic [ref=s1e267]:\\n - generic [ref=s1e269]\\n - generic [ref=s1e270]:\\n - text: Mapcarta\\n - generic [ref=s1e273]: https://mapcarta.com\\n - heading \"Building 99 Map - King County, Washington, USA - Mapcarta\" [level=2] [ref=s1e275]:\\n - link \"Building 99 Map - King County, Washington, USA - Mapcarta\" [ref=s1e276]:\\n - /url: https://mapcarta.com/W93639217\\n - paragraph [ref=s1e278]: Building 99 is a building in King County, Puget Sound,\\n Washington which is located on Northeast 36th Street. Building 99\\n is situated nearby to the food court Microsoft Cafe 99 , as well\\n as near …\\n - listitem [ref=s1e279]:\\n - link \"Microsoft\" [ref=s1e281]:\\n - /url: https://www.microsoft.com/en-us/about/office-locations\\n - generic [ref=s1e283]:\\n - generic [ref=s1e285]\\n - generic [ref=s1e286]:\\n - text: Microsoft\\n - generic [ref=s1e289]: https://www.microsoft.com › en-us › about …\\n - generic [ref=s1e292]:\\n - generic [ref=s1e294]:\\n - link \"/images/search?view=detailV2&ccid=REaQfaXR&id=A8AF32107EC802C72A6A5E9DBAA47E48A3D34B0C&thid=OIP.REaQfaXRYleauLHxKbPYGQAAAA&mediaurl=https://cdn-dynmedia-1.microsoft.com/is/image/microsoftcorp/About-OfficeLocations–OutdoorRedmond30-32-6484x4323&q=Microsoft+Building+99+address&ck=6C580933DF439BBF1D848AAC10172FA1&idpp=rc&idpview=singleimage&form=rc2idp&mode=overlay\" [ref=s1e295]:\\n - /url: javascript:void(0)\\n - heading \"Microsoft Office Locations | About Microsoft\" [level=2] [ref=s1e297]:\\n - link \"Microsoft Office Locations | About Microsoft\" [ref=s1e298]:\\n - /url: https://www.microsoft.com/en-us/about/office-locations\\n - paragraph [ref=s1e299]: Microsoft is based in Redmond, Washington with offices\\n across the US. Learn more about these locations. Microsoft’s\\n global headquarters are located on 500 acres in Redmond,\\n Washington that includes public spaces, sports fields, …\\n - generic [ref=s1e303]:\\n - heading \"Videos of Microsoft Building 99 Address\" [level=2] [ref=s1e305]:\\n - link \"Videos of Microsoft Building 99 Address\" [ref=s1e306]:\\n - /url: /videos/search?q=Microsoft+Building+99+address&qpvt=Microsoft+Building+99+address&FORM=VDRE\\n - generic [ref=s1e308]:\\n - generic [ref=s1e310]: bing.com › videos\\n - generic [ref=s1e312]:\\n - generic [ref=s1e314]:\\n - \\'link \"Interview and Q&A with Jenny Sabin, Creator of the Ada Installation in \\ue000Microsoft\\ue001 \\ue000Building\\ue001 \\ue00099\\ue001 from YouTube · Duration: 22 minutes 53 seconds · 1.4K views · uploaded on Oct 25, 2021 · uploaded by Microsoft Research · Click to play.\" [ref=s1e315]\\':\\n - /url: https://www.youtube.com/watch?v=BtgiDwS7w84\\n - generic [ref=s1e316]:\\n - generic [ref=s1e317]:\\n - img \"Interview and Q&A with Jenny Sabin, Creator of the\\n Ada Installation in Microsoft Building 99\" [ref=s1e319]\\n - generic [ref=s1e323]:\\n - generic [ref=s1e325]: 22:53\\n - generic [ref=s1e326]:\\n - generic \"Interview and Q&A with Jenny Sabin, Creator of the Ada Installation in Microsoft Building 99\" [ref=s1e327]:\\n - text: Interview and Q&A with Jenny Sabin, Creator of the Ada Installation in\\n - strong [ref=s1e328]: Microsoft\\n - strong [ref=s1e329]: Building\\n - strong [ref=s1e330]: \"99\"\\n - generic [ref=s1e331]:\\n - generic [ref=s1e332]:\\n - text: 1.4K views\\n - text: · Oct 25, 2021\\n - generic [ref=s1e335]:\\n - text: YouTube\\n - text: › Microsoft Research\\n - generic [ref=s1e339]:\\n - \\'link \"Inside \\ue000Microsoft\\ue001\\'\\'s Multi-Billion Dollar Headquarter from YouTube · Duration: 9 minutes 5 seconds · 3.9K views · uploaded on Jul 24, 2023 · uploaded by Lavish Woo · Click to play.\" [ref=s1e340]\\':\\n - /url: /videos/riverview/relatedvideo?q=Microsoft+Building+99+address&mid=609C51D0DBD3F35EA148609C51D0DBD3F35EA148&FORM=VIRE\\n - generic [ref=s1e341]:\\n - generic [ref=s1e342]:\\n - img \"Inside Microsoft\\'s Multi-Billion Dollar Headquarter\"\\n [ref=s1e344]\\n - generic [ref=s1e348]:\\n - generic [ref=s1e350]: 9:05\\n - generic [ref=s1e351]:\\n - generic \"Inside Microsoft\\'s Multi-Billion Dollar Headquarter\" [ref=s1e352]:\\n - text: Inside\\n - strong [ref=s1e353]: Microsoft\\n - text: \"\\'s Multi-Billion Dollar Headquarter\"\\n - generic [ref=s1e354]:\\n - generic [ref=s1e355]:\\n - text: 3.9K views\\n - text: · Jul 24, 2023\\n - generic [ref=s1e358]:\\n - text: YouTube\\n - text: › Lavish Woo\\n - generic [ref=s1e362]:\\n - \\'link \"Inside \\ue000Microsoft\\ue001\\'\\'s Insane Headquarters from YouTube · Duration: 10 minutes 15 seconds · 9K views · uploaded on Feb 27, 2022 · uploaded by Simply Tech · Click to play.\" [ref=s1e363]\\':\\n - /url: /videos/riverview/relatedvideo?q=Microsoft+Building+99+address&mid=B2D9E016C5679BA025F2B2D9E016C5679BA025F2&FORM=VIRE\\n - generic [ref=s1e364]:\\n - generic [ref=s1e365]:\\n - img \"Inside Microsoft\\'s Insane Headquarters\" [ref=s1e367]\\n - generic [ref=s1e371]:\\n - generic [ref=s1e373]: 10:15\\n - generic [ref=s1e374]:\\n - generic \"Inside Microsoft\\'s Insane Headquarters\" [ref=s1e375]:\\n - text: Inside\\n - strong [ref=s1e376]: Microsoft\\n - text: \"\\'s Insane Headquarters\"\\n - generic [ref=s1e377]:\\n - generic [ref=s1e378]:\\n - text: 9K views\\n - text: · Feb 27, 2022\\n - generic [ref=s1e381]:\\n - text: YouTube\\n - text: › Simply Tech\\n - generic [ref=s1e385]:\\n - \\'link \"Look Inside \\ue000Microsoft\\ue001\\'\\'s Massive Headquarters from YouTube · Duration: 4 minutes 20 seconds · 1.3K views · uploaded on Nov 27, 2022 · uploaded by Futurostructure - Infrastructure Of The Future · Click to play.\" [ref=s1e386]\\':\\n - /url: /videos/riverview/relatedvideo?q=Microsoft+Building+99+address&mid=761115F09388C2C34DB2761115F09388C2C34DB2&FORM=VIRE\\n - generic [ref=s1e387]:\\n - generic [ref=s1e388]:\\n - generic [ref=s1e390]\\n - generic [ref=s1e394]:\\n - generic [ref=s1e396]: 4:20\\n - generic [ref=s1e397]:\\n - generic \"Look Inside Microsoft\\'s Massive Headquarters\" [ref=s1e398]:\\n - text: Look Inside\\n - strong [ref=s1e399]: Microsoft\\n - text: \"\\'s Massive Headquarters\"\\n - generic [ref=s1e400]:\\n - generic [ref=s1e401]:\\n - text: 1.3K views\\n - text: · Nov 27, 2022\\n - generic [ref=s1e404]:\\n - text: YouTube\\n - text: › Futurostructure - Infrastructure Of The Future\\n - listitem [ref=s1e407]:\\n - link \"AES | Audio Engineering Society\" [ref=s1e409]:\\n - /url: https://www.aes.org/sections/pnw/direct/ms_rsch.htm\\n - generic [ref=s1e411]:\\n - generic [ref=s1e413]\\n - generic [ref=s1e414]:\\n - text: AES | Audio Engineering Society\\n - generic [ref=s1e417]: https://www.aes.org › sections › pnw › direct › ms_rsch.htm\\n - heading \"Directions to Microsoft Research\" [level=2] [ref=s1e419]:\\n - link \"Directions to Microsoft Research\":\\n - /url: https://www.aes.org/sections/pnw/direct/ms_rsch.htm\\n - paragraph [ref=s1e422]:\\n - text: Oct 9, 2018\\n - text: · Microsoft Research is located in Redmond, at the intersection of NE 36th\\n Street and 148th Avenue NE. This is south of where Microsoft\\n Studios are located. Microsoft Building 99\\n - listitem [ref=s1e424]:\\n - link \"campusbuilding.com\" [ref=s1e426]:\\n - /url: https://campusbuilding.com/company/microsoft/\\n - generic [ref=s1e428]:\\n - generic [ref=s1e430]\\n - generic [ref=s1e431]:\\n - text: campusbuilding.com\\n - generic [ref=s1e434]: https://campusbuilding.com › company › m…\\n - generic [ref=s1e437]:\\n - generic [ref=s1e439]:\\n - link \"/images/search?view=detailV2&ccid=X/EPBTCi&id=CF9C24AC67A42A5C771971C522BF3F3A97CC6655&thid=OIP.X_EPBTCi30m94XHqT52lnwHaH3&mediaurl=https://campusbuilding.com/static/images/seattle_area_microsoft_buildings_map.jpg&q=Microsoft+Building+99+address&ck=B3AEC9A066B8C9D9A861D34C2D56B0FD&idpp=rc&idpview=singleimage&form=rc2idp&mode=overlay\" [ref=s1e440]:\\n - /url: javascript:void(0)\\n - heading \"Microsoft Corporate Locations and Headquarters\" [level=2] [ref=s1e442]:\\n - link \"Microsoft Corporate Locations and Headquarters\" [ref=s1e443]:\\n - /url: https://campusbuilding.com/company/microsoft/\\n - paragraph [ref=s1e444]: Microsoft Corporate Locations and Headquarters In the\\n Seattle Area, Microsoft has 6 campuses and 132 buildings. There\\n have been 49 jobs posted in the last week.\\n - listitem [ref=s1e446]:\\n - link \"MapQuest\" [ref=s1e448]:\\n - /url: https://www.mapquest.com/us/washington/microsoft-building-99-parking-garage-472010688\\n - generic [ref=s1e450]:\\n - generic [ref=s1e452]\\n - generic [ref=s1e453]:\\n - text: MapQuest\\n - generic [ref=s1e456]: https://www.mapquest.com › us › washington\\n - heading \"Microsoft Building 99 Parking Garage - Official MapQuest\" [level=2] [ref=s1e458]:\\n - link \"Microsoft Building 99 Parking Garage - Official MapQuest\":\\n - /url: https://www.mapquest.com/us/washington/microsoft-building-99-parking-garage-472010688\\n - paragraph [ref=s1e461]: Microsoft Building 99 Parking Garage in Redmond, WA\\n offers convenient parking services for employees and visitors of\\n the company. The facility provides a secure and accessible\\n location …\\n - listitem [ref=s1e462]:\\n - link \"Place Digger\" [ref=s1e464]:\\n - /url: https://us.placedigger.com/microsoft-headquarters-and-visitor-center-redmon-seattle---usa27266365.html\\n - generic [ref=s1e466]:\\n - generic [ref=s1e468]\\n - generic [ref=s1e469]:\\n - text: Place Digger\\n - generic [ref=s1e472]: https://us.placedigger.com › microsoft-headquarters...\\n - heading \"Microsoft Headquarters & Visitor Center, Redmon (Seattle - U.S.A)\" [level=2] [ref=s1e474]:\\n - link \"Microsoft Headquarters & Visitor Center, Redmon (Seattle - U.S.A)\":\\n - /url: https://us.placedigger.com/microsoft-headquarters-and-visitor-center-redmon-seattle---usa27266365.html\\n - paragraph [ref=s1e477]: Microsoft Headquarters & Visitor Center, Redmon (Seattle\\n - U.S.A) is one of the popular Shopping & Retail located in 15010\\n NE 36th Street, Building 92 ,Redmond listed under Corporate Office\\n …\\n - listitem [ref=s1e478]:\\n - link \"cityseeker\" [ref=s1e480]:\\n - /url: https://cityseeker.com/redmond-wa/735585-microsoft-building-99\\n - generic [ref=s1e482]:\\n - generic [ref=s1e484]\\n - generic [ref=s1e485]:\\n - text: cityseeker\\n - generic [ref=s1e488]: https://cityseeker.com › redmond-wa\\n - heading \"Microsoft Building 99, Redmond - cityseeker\" [level=2] [ref=s1e490]:\\n - link \"Microsoft Building 99, Redmond - cityseeker\":\\n - /url: https://cityseeker.com/redmond-wa/735585-microsoft-building-99\\n - paragraph [ref=s1e493]: 14820 North East 36th Street, Microsoft Research Campus,\\n Redmond, WA, United States, 98052\\n - listitem [ref=s1e494]:\\n - generic [ref=s1e495]:\\n - heading \"Related searches for Microsoft Building 99 address\" [level=2] [ref=s1e496]:\\n - text: Related searches for\\n - strong [ref=s1e497]: Microsoft Building 99 address\\n - list [ref=s1e498]:\\n - listitem [ref=s1e499]:\\n - link \"microsoft 99 redmond\" [ref=s1e500]:\\n - /url: /search?q=microsoft+99+redmond&FORM=QSRE1\\n - generic [ref=s1e502]:\\n - text: microsoft 99\\n - strong [ref=s1e503]: redmond\\n - listitem [ref=s1e504]:\\n - link \"microsoft building 99 parking garage\" [ref=s1e505]:\\n - /url: /search?q=microsoft+building+99+parking+garage&FORM=QSRE2\\n - generic [ref=s1e507]:\\n - text: microsoft building 99\\n - strong [ref=s1e508]: parking garage\\n - listitem [ref=s1e509]:\\n - link \"inside microsoft headquarters\" [ref=s1e510]:\\n - /url: /search?q=inside+microsoft+headquarters&FORM=QSRE3\\n - generic [ref=s1e512]:\\n - strong [ref=s1e513]: inside\\n - text: microsoft\\n - strong [ref=s1e514]: headquarters\\n - listitem [ref=s1e515]:\\n - link \"microsoft anechoic chamber visit\" [ref=s1e516]:\\n - /url: /search?q=microsoft+anechoic+chamber+visit&FORM=QSRE4\\n - generic [ref=s1e518]:\\n - text: microsoft\\n - strong [ref=s1e519]: anechoic chamber visit\\n - listitem [ref=s1e520]:\\n - link \"microsoft redmond wa 98052\" [ref=s1e521]:\\n - /url: /search?q=microsoft+redmond+wa+98052&FORM=QSRE5\\n - generic [ref=s1e523]:\\n - text: microsoft\\n - strong [ref=s1e524]: redmond wa 98052\\n - listitem [ref=s1e525]:\\n - link \"microsoft anechoic chamber\" [ref=s1e526]:\\n - /url: /search?q=microsoft+anechoic+chamber&FORM=QSRE6\\n - generic [ref=s1e528]:\\n - text: microsoft\\n - strong [ref=s1e529]: anechoic chamber\\n - listitem [ref=s1e530]:\\n - link \"microsoft research redmond\" [ref=s1e531]:\\n - /url: /search?q=microsoft+research+redmond&FORM=QSRE7\\n - generic [ref=s1e533]:\\n - text: microsoft\\n - strong [ref=s1e534]: research redmond\\n - listitem [ref=s1e535]:\\n - link \"main microsoft campus\" [ref=s1e536]:\\n - /url: /search?q=main+microsoft+campus&FORM=QSRE8\\n - generic [ref=s1e538]:\\n - strong [ref=s1e539]: main\\n - text: microsoft\\n - strong [ref=s1e540]: campus\\n - listitem [ref=s1e541]:\\n - navigation \"More results for Microsoft Building 99 address\":\\n - list:\\n - listitem [ref=s1e544]: \"1\"\\n - listitem [ref=s1e546]:\\n - link \"Page 2\" [ref=s1e547]:\\n - /url: /search?q=Microsoft+Building+99+address&FPIG=629D5CE705334C83937EBCDDD6C544D1&first=11&FORM=PERE\\n - text: \"2\"\\n - listitem [ref=s1e548]:\\n - link \"Page 3\" [ref=s1e549]:\\n - /url: /search?q=Microsoft+Building+99+address&FPIG=629D5CE705334C83937EBCDDD6C544D1&first=21&FORM=PERE1\\n - text: \"3\"\\n - listitem [ref=s1e550]:\\n - link \"Page 4\" [ref=s1e551]:\\n - /url: /search?q=Microsoft+Building+99+address&FPIG=629D5CE705334C83937EBCDDD6C544D1&first=31&FORM=PERE2\\n - text: \"4\"\\n - listitem [ref=s1e552]:\\n - link \"Next page\" [ref=s1e553]:\\n - /url: /search?q=Microsoft+Building+99+address&FPIG=629D5CE705334C83937EBCDDD6C544D1&first=11&FORM=PORE\\n - complementary \"Additional Results\" [ref=s1e554]:\\n - list [ref=s1e555]:\\n - listitem [ref=s1e556]:\\n - generic [ref=s1e557]:\\n - generic [ref=s1e559]:\\n - generic [ref=s1e560]:\\n - generic [ref=s1e562]:\\n - heading \"Microsoft Building 99\" [level=2] [ref=s1e564]:\\n - link \"Microsoft Building 99\" [ref=s1e565]:\\n - /url: https://www.bing.com/alink/link?url=https%3a%2f%2fwww.microsoft.com%2f&source=serp-local&h=k6XBdzEhm26dMOehxO4ANkPmLgfNzfJEHe2c3sGHZUI%3d&p=lw_tpt&ig=629D5CE705334C83937EBCDDD6C544D1&ypid=YN873x101353856\\n - generic [ref=s1e566]:\\n - button \"Save\" [ref=s1e568]\\n - generic [ref=s1e570]:\\n - button \"Share\" [ref=s1e571]\\n - generic [ref=s1e574]: Software development in Redmond, Wa\\n - generic [ref=s1e576]:\\n - generic [ref=s1e578]:\\n - generic [ref=s1e579]:\\n - link \"Website\" [ref=s1e580]:\\n - /url: https://www.bing.com/alink/link?url=https%3a%2f%2fwww.microsoft.com%2f&source=serp-local&h=k6XBdzEhm26dMOehxO4ANkPmLgfNzfJEHe2c3sGHZUI%3d&p=lw_tp&ig=629D5CE705334C83937EBCDDD6C544D1&ypid=YN873x101353856\\n - img [ref=s1e582]\\n - link \"Directions\" [ref=s1e586]:\\n - /url: /maps?&mepi=127~Directions~Unknown~Direction_Button&ty=0&rtp=pos.47.64213943481445_-122.14218139648438__Microsoft%20Building%2099__e_~&mode=d&v=2&sV=1\\n - img [ref=s1e588]\\n - list [ref=s1e592]:\\n - listitem [ref=s1e593]:\\n - button \"Prices\" [ref=s1e594]\\n - generic [ref=s1e596]:\\n - group \"Address\" [ref=s1e597]:\\n - img [ref=s1e598]\\n - generic [ref=s1e602]:\\n - link \"14820 NE 36th St, Redmond, Wa 98052\" [ref=s1e604]:\\n - /url: /maps?&mepi=127~~Unknown~Address_Link&ty=18&q=Microsoft%20Building%2099&ss=ypid.YN873x101353856&ppois=47.64213943481445_-122.14218139648438_Microsoft%20Building%2099_YN873x101353856~&cp=47.642139~-122.142181&v=2&sV=1&FORM=MPSRPL\\n - text: · 1.1 mi\\n - group \"Phone\" [ref=s1e605]:\\n - img [ref=s1e606]\\n - link \"Phone (425) 882-8080\" [ref=s1e610]:\\n - /url: tel:4258828080\\n - text: (425) 882-8080\\n - generic [ref=s1e611]:\\n - img [ref=s1e612]\\n - generic [ref=s1e615]:\\n - button \"Suggest an edit\" [ref=s1e616]\\n - text: ·\\n - text: Your business?\\n - link \"Claim now\" [ref=s1e618]:\\n - /url: https://www.bingplaces.com/DashBoard/Edit?Id=YN873x101353856&market=en-US&src=SERPIC\\n - generic [ref=s1e620]:\\n - heading \"Add more information\" [level=2] [ref=s1e621]\\n - generic [ref=s1e623]:\\n - generic [ref=s1e624]:\\n - img [ref=s1e625]\\n - button \"Add hours\" [ref=s1e629]\\n - generic [ref=s1e631]:\\n - generic [ref=s1e633]:\\n - button \"Add photos\" [ref=s1e634]:\\n - img [ref=s1e635]\\n - text: Add photos\\n - generic [ref=s1e640]:\\n - text: Microsoft creates platforms and tools powered by AI to deliver innovative\\n solutions that meet the evolving needs of our customers.\\n The technology …\\n - link \"See more See more\" [ref=s1e642]:\\n - /url: https://news.microsoft.com/facts-about-microsoft\\n - text: See more\\n - generic [ref=s1e644]:\\n - heading \"Frequently asked questions\" [level=2] [ref=s1e645]\\n - generic [ref=s1e646]:\\n - generic [ref=s1e647]:\\n - generic [ref=s1e649]:\\n - text: \"Q:\"\\n - generic [ref=s1e652]: What is the difference between Microsoft 365\\n (subscription) and Office 2024 (one-time purchase)?\\n - generic [ref=s1e654]:\\n - generic [ref=s1e655]:\\n - text: \"A:\"\\n - generic [ref=s1e658]:\\n - text: Microsoft 365 is a subscription that includes the most collaborative,\\n up-to-date features in one seamless, integrated\\n experience. Microsoft 365 includes the …\\n - button \"Show more\" [ref=s1e660]\\n - generic [ref=s1e663]:\\n - link \"Read more\" [ref=s1e665]:\\n - /url: https://microsoft.com/en-us/microsoft-365/microsoft-365-for-home-and-school-faq\\n - link \"See all 50 questions\" [ref=s1e667]:\\n - /url: \"#\"\\n - generic [ref=s1e668]:\\n - text: \"Data from:\"\\n - link \"BusinessWebsite\" [ref=s1e669]:\\n - /url: https://microsoft.com/en-us/microsoft-365/microsoft-365-for-home-and-school-faq\\n - generic [ref=s1e672]:\\n - heading \"Social profiles\" [level=2] [ref=s1e673]\\n - group \"Social profiles\" [ref=s1e675]:\\n - list [ref=s1e676]:\\n - listitem [ref=s1e677]:\\n - link \"Facebook icon Facebook\" [ref=s1e678]:\\n - /url: https://www.facebook.com/Microsoft\\n - img \"Facebook icon\" [ref=s1e680]\\n - text: Facebook\\n - listitem [ref=s1e682]:\\n - link \"X icon X\" [ref=s1e683]:\\n - /url: https://twitter.com/microsoft\\n - img \"X icon\" [ref=s1e685]\\n - text: X\\n - listitem [ref=s1e687]:\\n - link \"LinkedIn icon LinkedIn\" [ref=s1e688]:\\n - /url: https://www.linkedin.com/company/microsoft\\n - img \"LinkedIn icon\" [ref=s1e690]\\n - text: LinkedIn\\n - generic [ref=s1e693]:\\n - heading \"People also search for\" [level=2] [ref=s1e694]\\n - generic [ref=s1e695]:\\n - text: Software development\\n - generic [ref=s1e698]:\\n - generic [ref=s1e700]:\\n - generic [ref=s1e702]:\\n - list \"Please use arrow keys to navigate\" [ref=s1e703]:\\n - listitem [ref=s1e704]:\\n - link \"Acumatica Cloud ERP Acumatica Cloud ERP\" [ref=s1e705]:\\n - /url: /search?q=Acumatica+Cloud+ERP&filters=local_ypid%3a%22873x11271913447243333430%22&FORM=SNAPST\\n - generic [ref=s1e706]:\\n - img \"Acumatica Cloud ERP\" [ref=s1e708]\\n - generic [ref=s1e710]: Acumatica Cloud ERP\\n - listitem [ref=s1e711]:\\n - link \"AscendoSoft Inc. AscendoSoft Inc.\" [ref=s1e712]:\\n - /url: /search?q=AscendoSoft+Inc.&filters=local_ypid%3a%22873x109094970%22&FORM=SNAPST\\n - generic [ref=s1e713]:\\n - img \"AscendoSoft Inc.\" [ref=s1e715]\\n - generic [ref=s1e717]: AscendoSoft Inc.\\n - listitem [ref=s1e718]:\\n - link \"TecAce Software, Ltd TecAce Software, Ltd\" [ref=s1e719]:\\n - /url: /search?q=TecAce+Software%2c+Ltd&filters=local_ypid%3a%22873x100939365%22&FORM=SNAPST\\n - generic [ref=s1e720]:\\n - img \"TecAce Software, Ltd\" [ref=s1e722]\\n - generic [ref=s1e724]: TecAce Software, Ltd\\n - listitem [ref=s1e725]:\\n - link \"Cirkled In\" [ref=s1e726]:\\n - /url: /search?q=Cirkled+In&filters=local_ypid%3a%22873x10666105865648718724%22&FORM=SNAPST\\n - generic [ref=s1e727]:\\n - generic [ref=s1e729]\\n - generic [ref=s1e731]: Cirkled In\\n - listitem [ref=s1e732]:\\n - link \"Vishwak Solutions Inc\" [ref=s1e733]:\\n - /url: /search?q=Vishwak+Solutions+Inc&filters=local_ypid%3a%22873x110407396%22&FORM=SNAPST\\n - generic [ref=s1e734]:\\n - generic [ref=s1e736]\\n - generic [ref=s1e738]: Vishwak Solutions Inc\\n - generic [ref=s1e739]:\\n - text: IT service & computer repair\\n - generic [ref=s1e742]:\\n - generic [ref=s1e744]:\\n - generic [ref=s1e746]:\\n - list \"Please use arrow keys to navigate\" [ref=s1e747]:\\n - listitem [ref=s1e748]:\\n - link \"Kirwan Computer\" [ref=s1e749]:\\n - /url: /search?q=Kirwan+Computer&filters=local_ypid%3a%22873x114284629%22&FORM=SNAPST\\n - generic [ref=s1e750]:\\n - generic [ref=s1e752]\\n - generic [ref=s1e754]: Kirwan Computer\\n - listitem [ref=s1e755]:\\n - link \"Digital forensics\" [ref=s1e756]:\\n - /url: /search?q=Digital+forensics&filters=local_ypid%3a%22873x13044985277069239607%22&FORM=SNAPST\\n - generic [ref=s1e757]:\\n - generic [ref=s1e759]\\n - generic [ref=s1e761]: Digital forensics\\n - listitem [ref=s1e763]:\\n - generic [ref=s1e765]:\\n - generic [ref=s1e766]:\\n - heading \"Related searches for Microsoft Building 99 address\" [level=2] [ref=s1e768]:\\n - text: Related searches for\\n - strong [ref=s1e769]: Microsoft Building 99 address\\n - generic [ref=s1e770]:\\n - link \"microsoft 99 redmond\" [ref=s1e772]:\\n - /url: /search?q=microsoft+99+redmond&FORM=R5FD\\n - generic [ref=s1e774]:\\n - text: microsoft 99\\n - strong [ref=s1e775]: redmond\\n - link \"microsoft building 99 parking garage\" [ref=s1e777]:\\n - /url: /search?q=microsoft+building+99+parking+garage&FORM=R5FD1\\n - generic [ref=s1e779]:\\n - text: microsoft building 99\\n - strong [ref=s1e780]: parking garage\\n - link \"inside microsoft headquarters\" [ref=s1e782]:\\n - /url: /search?q=inside+microsoft+headquarters&FORM=R5FD2\\n - generic [ref=s1e784]:\\n - strong [ref=s1e785]: inside\\n - text: microsoft\\n - strong [ref=s1e786]: headquarters\\n - link \"microsoft anechoic chamber visit\" [ref=s1e788]:\\n - /url: /search?q=microsoft+anechoic+chamber+visit&FORM=R5FD3\\n - generic [ref=s1e790]:\\n - text: microsoft\\n - strong [ref=s1e791]: anechoic chamber visit\\n - link \"microsoft redmond wa 98052\" [ref=s1e793]:\\n - /url: /search?q=microsoft+redmond+wa+98052&FORM=R5FD4\\n - generic [ref=s1e795]:\\n - text: microsoft\\n - strong [ref=s1e796]: redmond wa 98052\\n - link \"microsoft anechoic chamber\" [ref=s1e798]:\\n - /url: /search?q=microsoft+anechoic+chamber&FORM=R5FD5\\n - generic [ref=s1e800]:\\n - text: microsoft\\n - strong [ref=s1e801]: anechoic chamber\\n - link \"microsoft research redmond\" [ref=s1e803]:\\n - /url: /search?q=microsoft+research+redmond&FORM=R5FD6\\n - generic [ref=s1e805]:\\n - text: microsoft\\n - strong [ref=s1e806]: research redmond\\n - link \"main microsoft campus\" [ref=s1e808]:\\n - /url: /search?q=main+microsoft+campus&FORM=R5FD7\\n - generic [ref=s1e810]:\\n - strong [ref=s1e811]: main\\n - text: microsoft\\n - strong [ref=s1e812]: campus\\n - button \"Feedback\" [ref=s1e813]\\n - contentinfo \"Footer\" [ref=s1e815]:\\n - generic [ref=s1e816]:\\n - text: © 2025 Microsoft\\n - list:\\n - listitem [ref=s1e819]:\\n - link \"Privacy and Cookies\" [ref=s1e820]:\\n - /url: http://go.microsoft.com/fwlink/?LinkId=521839\\n - listitem [ref=s1e821]:\\n - link \"Legal\" [ref=s1e822]:\\n - /url: http://go.microsoft.com/fwlink/?LinkID=246338\\n - listitem [ref=s1e823]:\\n - link \"Advertise\" [ref=s1e824]:\\n - /url: https://go.microsoft.com/fwlink/?linkid=868922\\n - listitem [ref=s1e825]:\\n - link \"About our ads\" [ref=s1e826]:\\n - /url: http://go.microsoft.com/fwlink/?LinkID=286759\\n - listitem [ref=s1e827]:\\n - link \"Help\" [ref=s1e828]:\\n - /url: https://support.microsoft.com/topic/82d20721-2d6f-4012-a13d-d1910ccf203f\\n - listitem [ref=s1e829]:\\n - button \"Feedback\" [ref=s1e830]\\n - listitem [ref=s1e831]:\\n - link \"Your Privacy Choices\" [ref=s1e832]:\\n - /url: https://go.microsoft.com/fwlink/?linkid=2214802\\n - listitem [ref=s1e833]:\\n - link \"Consumer Health Privacy\" [ref=s1e834]:\\n - /url: https://go.microsoft.com/fwlink/?linkid=2259814\\n - link \"🐞\" [ref=s1e835]:\\n - /url: javascript:void(0)\\n```')] is_error=False\n", + "---------Final Response-----------\n", + "The address of Microsoft Building 99 is 14820 NE 36th Street, Redmond, WA 98052, United States.\n" + ] + } + ], + "source": [ + "from autogen_core import AgentId, SingleThreadedAgentRuntime\n", + "from autogen_core.model_context import BufferedChatCompletionContext\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "from autogen_ext.tools.mcp import McpWorkbench, SseServerParams\n", + "\n", + "playwright_server_params = SseServerParams(\n", + " url=\"http://localhost:8931/sse\",\n", + ")\n", + "\n", + "# Start the workbench in a context manager.\n", + "# You can also start and stop the workbench using `workbench.start()` and `workbench.stop()`.\n", + "async with McpWorkbench(playwright_server_params) as workbench: # type: ignore\n", + " # Create a single-threaded agent runtime.\n", + " runtime = SingleThreadedAgentRuntime()\n", + "\n", + " # Register the agent with the runtime.\n", + " await WorkbenchAgent.register(\n", + " runtime=runtime,\n", + " type=\"WebAgent\",\n", + " factory=lambda: WorkbenchAgent(\n", + " model_client=OpenAIChatCompletionClient(model=\"gpt-4.1-nano\"),\n", + " model_context=BufferedChatCompletionContext(buffer_size=10),\n", + " workbench=workbench,\n", + " ),\n", + " )\n", + "\n", + " # Start the runtime.\n", + " runtime.start()\n", + "\n", + " # Send a message to the agent.\n", + " await runtime.send_message(\n", + " Message(content=\"Use Bing to find out the address of Microsoft Building 99\"),\n", + " recipient=AgentId(\"WebAgent\", \"default\"),\n", + " )\n", + "\n", + " # Stop the runtime.\n", + " await runtime.stop()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/azure-openai-with-aad-auth.md b/python/docs/src/user-guide/core-user-guide/cookbook/azure-openai-with-aad-auth.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/azure-openai-with-aad-auth.md rename to python/docs/src/user-guide/core-user-guide/cookbook/azure-openai-with-aad-auth.md diff --git a/python/docs/src/user-guide/core-user-guide/cookbook/data/nifty_500_quarterly_results.csv b/python/docs/src/user-guide/core-user-guide/cookbook/data/nifty_500_quarterly_results.csv new file mode 100644 index 000000000000..e02068e09042 --- /dev/null +++ b/python/docs/src/user-guide/core-user-guide/cookbook/data/nifty_500_quarterly_results.csv @@ -0,0 +1,502 @@ +name,NSE_code,BSE_code,sector,industry,revenue,operating_expenses,operating_profit,operating_profit_margin,depreciation,interest,profit_before_tax,tax,net_profit,EPS,profit_TTM,EPS_TTM +3M India Ltd.,3MINDIA,523395,GENERAL INDUSTRIALS,INDUSTRIAL MACHINERY,"1,057",847.4,192.1,18.48%,12.9,0.7,195.9,49.8,146.1,129.7,535.9,475.7 +ACC Ltd.,ACC,500410,CEMENT AND CONSTRUCTION,CEMENT & CEMENT PRODUCTS,"4,644.8","3,885.4",549.3,12.39%,212.8,28.9,517.7,131.5,387.9,20.7,"1,202.7",64 +AIA Engineering Ltd.,AIAENG,532683,GENERAL INDUSTRIALS,OTHER INDUSTRIAL GOODS,"1,357.1",912.7,382.1,29.51%,24.5,7.4,412.5,88.4,323.1,34.3,"1,216.1",128.9 +APL Apollo Tubes Ltd.,APLAPOLLO,533758,METALS & MINING,IRON & STEEL PRODUCTS,"4,65","4,305.4",325,7.02%,41.3,26.6,276.7,73.8,202.9,7.3,767.5,27.7 +Au Small Finance Bank Ltd.,AUBANK,540611,BANKING AND FINANCE,BANKS,"2,956.5","1,026.7",647.7,25.59%,0,"1,282.1",533.4,131.5,401.8,6,"1,606.2",24 +Adani Ports & Special Economic Zone Ltd.,ADANIPORTS,532921,TRANSPORTATION,MARINE PORT & SERVICES,"6,951.9","2,982.4","3,664",55.13%,974.5,520.1,"2,474.9",759,"1,747.8",8.1,"6,337",29.3 +Adani Energy Solutions Ltd.,ADANIENSOL,ASM,UTILITIES,ELECTRIC UTILITIES,"3,766.5","2,169.3","1,504.6",40.95%,432.1,640.8,369.9,84.9,275.9,2.5,"1,315.1",11.8 +Aditya Birla Fashion and Retail Ltd.,ABFRL,535755,RETAILING,DEPARTMENT STORES,"3,272.2","2,903.6",322.9,10.01%,388.8,208.4,-228.6,-28.2,-179.2,-1.9,-491.7,-5.2 +Aegis Logistics Ltd.,AEGISCHEM,500003,OIL & GAS,OIL MARKETING & DISTRIBUTION,"1,279.3","1,026.5",208.3,16.87%,34.1,26.6,192,42,127,3.6,509,14.5 +Ajanta Pharma Ltd.,AJANTPHARM,532331,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"1,049.8",737.8,290.7,28.26%,33.7,2.3,275.9,80.6,195.3,15.5,660.2,52.3 +Alembic Pharmaceuticals Ltd.,APLLTD,533573,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"1,605.1","1,386.7",208.2,13.06%,67.6,15.7,135.1,-1.9,136.6,7,531.7,27 +Alkem Laboratories Ltd.,ALKEM,539523,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"3,503.4","2,693.4",746.7,21.71%,73.9,30.3,648,33.1,620.5,51.9,"1,432.9",119.9 +Amara Raja Energy & Mobility Ltd.,ARE&M,500008,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"2,988.6","2,556.9",402.5,13.60%,115.7,6.2,309.8,83.5,226.3,13.2,779.8,45.7 +Ambuja Cements Ltd.,AMBUJACEM,500425,CEMENT AND CONSTRUCTION,CEMENT & CEMENT PRODUCTS,"7,9","6,122.1","1,301.8",17.54%,380.9,61.2,"1,335.7",352.5,793,4,"2,777.9",14 +Apollo Hospitals Enterprise Ltd.,APOLLOHOSP,508869,DIVERSIFIED CONSUMER SERVICES,HEALTHCARE FACILITIES,"4,869.1","4,219.4",627.5,12.95%,163.4,111.3,376.9,130.2,232.9,16.2,697.5,48.5 +Apollo Tyres Ltd.,APOLLOTYRE,500877,AUTOMOBILES & AUTO COMPONENTS,AUTO TYRES & RUBBER PRODUCTS,"6,304.9","5,119.8","1,159.8",18.47%,360.3,132.8,679.9,205.8,474.3,7.5,"1,590.7",25 +Ashok Leyland Ltd.,ASHOKLEY,500477,AUTOMOBILES & AUTO COMPONENTS,COMMERCIAL VEHICLES,"11,463","9,558.6","1,870.4",16.37%,226.6,715.1,924.4,358,526,1.8,"2,141.5",7.3 +Asian Paints Ltd.,ASIANPAINT,500820,DIVERSIFIED CONSUMER SERVICES,FURNITURE-FURNISHING-PAINTS,"8,643.8","6,762.3","1,716.2",20.24%,208.7,50.9,"1,621.8",418.6,"1,205.4",12.6,"5,062.6",52.8 +Astral Ltd.,ASTRAL,532830,GENERAL INDUSTRIALS,PLASTIC PRODUCTS,"1,376.4","1,142.9",220.1,16.15%,48.7,8,176.8,45.1,131.2,4.9,549.7,20.4 +Atul Ltd.,ATUL,500027,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,"1,215.8","1,038.5",155.2,13.00%,54,1.9,121.5,32.5,90.3,30.6,392.3,132.9 +Aurobindo Pharma Ltd.,AUROPHARMA,524804,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"7,406.4","5,846","1,373.4",19.02%,417.5,68.2,"1,074.7",323.7,757.2,12.8,"2,325.5",39.7 +Avanti Feeds Ltd.,AVANTIFEED,512573,FOOD BEVERAGES & TOBACCO,OTHER FOOD PRODUCTS,"1,312","1,184.5",94,7.35%,14.3,0.2,113,30.5,74.2,5.5,336.4,24.7 +Avenue Supermarts Ltd.,DMART,540376,RETAILING,DEPARTMENT STORES,"12,661.3","11,619.4","1,005",7.96%,174.4,15.6,851.9,228.6,623.6,9.6,"2,332.1",35.8 +Axis Bank Ltd.,AXISBANK,532215,BANKING AND FINANCE,BANKS,"33,122.2","9,207.3","9,166",33.43%,0,"14,749","8,313.8","2,096.1","6,204.1",20.1,"13,121",42.6 +Bajaj Auto Ltd.,BAJAJ-AUTO,532977,AUTOMOBILES & AUTO COMPONENTS,2/3 WHEELERS,"11,206.8","8,708.1","2,130.1",19.65%,91.8,6.5,"2,400.4",564,"2,02",71.4,"6,841.6",241.8 +Bajaj Finance Ltd.,BAJFINANCE,500034,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),"13,381.8","2,851.5","9,449.7",70.63%,158.5,"4,537.1","4,757.6","1,207","3,550.8",58.7,"13,118.5",216.7 +Bajaj Finserv Ltd.,BAJAJFINSV,532978,DIVERSIFIED,HOLDING COMPANIES,"26,022.7","14,992.2","9,949.9",38.24%,208.8,"4,449.1","5,292","1,536.5","1,929",12.1,"7,422.6",46.6 +Bajaj Holdings & Investment Ltd.,BAJAJHLDNG,500490,DIVERSIFIED,HOLDING COMPANIES,240.1,33.5,191.2,85.08%,8.4,0.2,197.9,73.9,"1,491.2",134,"5,545.1",498.3 +Balkrishna Industries Ltd.,BALKRISIND,502355,AUTOMOBILES & AUTO COMPONENTS,AUTO TYRES & RUBBER PRODUCTS,"2,360.3","1,720.5",532.7,23.64%,160.4,23.9,455.5,108.1,347.4,18,"1,047.5",54.2 +Balrampur Chini Mills Ltd.,BALRAMCHIN,500038,FOOD BEVERAGES & TOBACCO,SUGAR,"1,649","1,374.6",164.9,10.71%,41.2,17.2,215.9,56.6,166.3,8.2,540.5,26.8 +Bank of Baroda,BANKBARODA,532134,BANKING AND FINANCE,BANKS,"35,766","8,430.4","9,807.9",33.52%,0,"17,527.7","6,022.8","1,679.7","4,458.4",8.5,"18,602.9",35.9 +Bank of India,BANKINDIA,532149,BANKING AND FINANCE,BANKS,"16,779.4","3,704.9","3,818.8",25.35%,0,"9,255.7","2,977.4","1,488.6","1,498.5",3.6,"5,388.7",13.1 +Bata India Ltd.,BATAINDIA,500043,RETAILING,FOOTWEAR,834.6,637.5,181.7,22.18%,81.7,28.4,46.1,12.1,34,2.6,289.7,22.5 +Berger Paints (India) Ltd.,BERGEPAINT,509480,DIVERSIFIED CONSUMER SERVICES,FURNITURE-FURNISHING-PAINTS,"2,782.6","2,293.7",473.6,17.12%,82.9,21.1,385,96.7,291.6,2.5,"1,032.6",8.9 +Bharat Electronics Ltd.,BEL,500049,GENERAL INDUSTRIALS,DEFENCE,"4,146.1","2,994.9","1,014.2",25.30%,108.3,1.5,"1,041.5",260.7,789.4,1.1,"3,323",4.5 +Bharat Forge Ltd.,BHARATFORG,500493,GENERAL INDUSTRIALS,OTHER INDUSTRIAL PRODUCTS,"3,826.7","3,152.8",621.4,16.47%,211.3,124.3,336.1,121.8,227.2,4.9,783.7,16.8 +Bharat Heavy Electricals Ltd.,BHEL,500103,GENERAL INDUSTRIALS,HEAVY ELECTRICAL EQUIPMENT,"5,305.4","5,513",-387.7,-7.56%,59.9,180.4,-447.9,-197.9,-238.1,-0.7,71.3,0.2 +Bharat Petroleum Corporation Ltd.,BPCL,500547,OIL & GAS,REFINERIES/PETRO-PRODUCTS,"103,72","90,103.9","12,940.5",12.56%,"1,605.3",973.2,"10,755.7","2,812.2","8,243.5",38.7,"27,505.3",129.2 +Bharti Airtel Ltd.,BHARTIARTL,532454,TELECOM SERVICES,TELECOM SERVICES,"37,374.2","17,530.1","19,513.7",52.68%,"9,734.3","5,185.8","3,353.7","1,846.5","1,340.7",2.4,"7,547",13.2 +Indus Towers Ltd.,INDUSTOWER,534816,TELECOM SERVICES,OTHER TELECOM SERVICES,"7,229.7","3,498.8","3,633.7",50.95%,"1,525.6",458.6,"1,746.7",452,"1,294.7",4.8,"3,333.5",12.4 +Biocon Ltd.,BIOCON,532523,PHARMACEUTICALS & BIOTECHNOLOGY,BIOTECHNOLOGY,"3,620.2","2,720.7",741.6,21.42%,389.3,247.7,238.5,41.6,125.6,1.1,498.4,4.2 +Birla Corporation Ltd.,BIRLACORPN,500335,CEMENT AND CONSTRUCTION,CEMENT & CEMENT PRODUCTS,"2,313.2","1,997",288.9,12.64%,143.5,95.4,77.1,18.8,58.4,7.6,153.1,19.9 +Blue Dart Express Ltd.,BLUEDART,526612,TRANSPORTATION,TRANSPORTATION - LOGISTICS,"1,329.7","1,101.8",222.7,16.82%,110.6,19.5,97.9,24.8,73.1,30.8,292.4,123.2 +Blue Star Ltd.,BLUESTARCO,500067,CONSUMER DURABLES,CONSUMER ELECTRONICS,"1,903.4","1,767.7",122.7,6.49%,23,17.6,95,24.3,70.7,3.6,437.7,21.3 +Bombay Burmah Trading Corporation Ltd.,BBTC,501425,FOOD BEVERAGES & TOBACCO,TEA & COFFEE,"4,643.5","3,664.7",859.2,18.99%,74.7,154.6,697.1,212.6,122,17.5,"-1,499.5",-214.8 +Bosch Ltd.,BOSCHLTD,500530,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"4,284.3","3,638.8",491.3,11.90%,101.3,12.2,"1,317",318.1,999.8,339,"2,126.9",721 +Brigade Enterprises Ltd.,BRIGADE,532929,REALTY,REALTY,"1,407.9","1,041.8",324.8,23.77%,75.7,110,180.3,67.8,133.5,5.8,298.2,12.9 +Britannia Industries Ltd.,BRITANNIA,500825,FMCG,PACKAGED FOODS,"4,485.2","3,560.5",872.4,19.68%,71.7,53.4,799.7,212.1,587.6,24.4,"2,536.2",105.3 +CCL Products India Ltd.,CCL,519600,FOOD BEVERAGES & TOBACCO,TEA & COFFEE,608.3,497.7,109.9,18.09%,22.6,18.4,69.7,8.8,60.9,4.6,279.9,21 +Crisil Ltd.,CRISIL,500092,BANKING AND FINANCE,OTHER FINANCIAL SERVICES,771.8,544.2,191.7,26.05%,26.5,0.8,200.3,48.3,152,20.8,606.3,82.9 +Zydus Lifesciences Ltd.,ZYDUSLIFE,532321,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"4,422.8","3,222.7","1,146.1",26.23%,184.2,8.7,"1,007.2",226.4,800.7,7.9,"2,807.1",27.7 +Can Fin Homes Ltd.,CANFINHOME,511196,BANKING AND FINANCE,HOUSING FINANCE,871,49.7,749.2,86.01%,2.8,548.4,198,39.9,158.1,11.9,658.8,49.5 +Canara Bank,CANBK,532483,BANKING AND FINANCE,BANKS,"33,891.2","8,250.3","7,706.6",28.24%,0,"17,934.3","5,098","1,420.6","3,86",20.9,"13,968.4",77 +Carborundum Universal Ltd.,CARBORUNIV,513375,GENERAL INDUSTRIALS,OTHER INDUSTRIAL PRODUCTS,"1,166",978.8,167.5,14.61%,45.9,4.9,136.4,43.7,101.9,5.4,461.3,24.3 +Castrol India Ltd.,CASTROLIND,500870,OIL & GAS,OIL MARKETING & DISTRIBUTION,"1,203.2",914.4,268.6,22.70%,22.9,2.4,263.5,69.1,194.4,2,815.5,8.2 +Ceat Ltd.,CEATLTD,500878,AUTOMOBILES & AUTO COMPONENTS,AUTO TYRES & RUBBER PRODUCTS,"3,063.8","2,597.2",456.1,14.94%,124.5,71.7,270.4,68.3,208,51.4,521.7,129 +Central Bank of India,CENTRALBK,532885,BANKING AND FINANCE,BANKS,"8,438.5","2,565.4","1,535.4",20.81%,0,"4,337.7",567.2,-41.5,622,0.7,"2,181.4",2.5 +Century Plyboards (India) Ltd.,CENTURYPLY,532548,FOREST MATERIALS,FOREST PRODUCTS,"1,011.4",852.5,144.3,14.47%,23.4,6.1,129.4,32.2,96.9,4.4,380.7,17.1 +Cera Sanitaryware Ltd.,CERA,532443,DIVERSIFIED CONSUMER SERVICES,FURNITURE-FURNISHING-PAINTS,476.2,387.2,76.5,16.49%,8.9,1.4,77.2,19.8,56.9,43.8,232.4,178.7 +Chambal Fertilisers & Chemicals Ltd.,CHAMBLFERT,500085,FERTILIZERS,FERTILIZERS,"5,467.3","4,770.5",615,11.42%,78.4,45.8,572.6,200.2,381,9.2,"1,137.7",27.3 +Cholamandalam Investment & Finance Company Ltd.,CHOLAFIN,511243,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),"4,695.2",987.6,"3,235.1",69.99%,38.5,"2,204.2","1,065",288.8,772.9,9.4,"3,022.8",36.7 +Cipla Ltd.,CIPLA,500087,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"6,854.5","4,944.4","1,733.8",25.96%,290,25.8,"1,594.2",438.4,"1,130.9",14,"3,449.1",42.7 +City Union Bank Ltd.,CUB,532210,BANKING AND FINANCE,BANKS,"1,486.1",333.9,386.6,29.65%,0,765.6,330.6,50,280.6,3.8,943.8,12.7 +Coal India Ltd.,COALINDIA,533278,METALS & MINING,COAL,"34,760.3","24,639.4","8,137",24.83%,"1,178.2",182.5,"8,760.2","2,036.5","6,799.8",11,"28,059.6",45.5 +Colgate-Palmolive (India) Ltd.,COLPAL,500830,FMCG,PERSONAL PRODUCTS,"1,492.1",989,482.1,32.77%,44.3,1.1,457.8,117.8,340.1,12.5,"1,173.2",43.1 +Container Corporation of India Ltd.,CONCOR,531344,COMMERCIAL SERVICES & SUPPLIES,WAREHOUSING AND LOGISTICS,"2,299.8","1,648.4",546.5,24.90%,153.1,16.5,481.8,119,367.4,6,"1,186.2",19.5 +Coromandel International Ltd.,COROMANDEL,506395,FERTILIZERS,FERTILIZERS,"7,032.9","5,929.4","1,058.7",15.15%,54,46.2,"1,003.3",245,756.9,25.7,"2,024.2",68.8 +Crompton Greaves Consumer Electricals Ltd.,CROMPTON,539876,CONSUMER DURABLES,HOUSEHOLD APPLIANCES,"1,797.2","1,607.8",174.5,9.79%,32.1,21.5,135.8,34.9,97.2,1.5,432,6.7 +Cummins India Ltd.,CUMMINSIND,500480,GENERAL INDUSTRIALS,INDUSTRIAL MACHINERY,"2,011.3","1,575.4",346.2,18.02%,38.3,6.8,390.9,99.6,329.1,11.9,"1,445.5",52.1 +Cyient Ltd.,CYIENT,532175,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"1,792","1,452.7",325.8,18.32%,65.8,27,240.3,56.7,178.3,16.3,665.6,60.1 +DCM Shriram Ltd.,DCMSHRIRAM,523367,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,"2,73","2,593.9",114.1,4.21%,74,14.7,47.5,15.2,32.2,2.1,617.6,39.4 +DLF Ltd.,DLF,532868,REALTY,REALTY,"1,476.4",885.3,462.4,34.31%,37,90.2,464,112.2,622.8,2.5,"2,239",9 +Dabur India Ltd.,DABUR,500096,FMCG,PERSONAL PRODUCTS,"3,320.2","2,543",660.9,20.63%,98.3,28.1,650.8,144.3,515,2.9,"1,755.7",9.9 +Delta Corp Ltd.,DELTACORP,532848,COMMERCIAL SERVICES & SUPPLIES,MISC. COMMERCIAL SERVICES,282.6,170.5,100.1,36.99%,16.9,2.7,92.4,23,69.4,2.6,273.3,10.2 +Divi's Laboratories Ltd.,DIVISLAB,532488,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"1,995","1,43",479,25.09%,95,1,469,121,348,13.1,"1,331.8",50.3 +Dr. Lal Pathlabs Ltd.,LALPATHLAB,539524,DIVERSIFIED CONSUMER SERVICES,HEALTHCARE SERVICES,619.4,423.5,177.8,29.57%,35.9,7.8,152.2,41.5,109.3,13.2,301.4,36.1 +Dr. Reddy's Laboratories Ltd.,DRREDDY,500124,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"7,217.6","4,888.8","2,008.3",29.09%,375.5,35.3,"1,912.5",434.5,"1,482.2",89.1,"5,091.2",305.2 +EID Parry (India) Ltd.,EIDPARRY,500125,FOOD BEVERAGES & TOBACCO,OTHER FOOD PRODUCTS,"9,210.3","8,002","1,057.5",11.67%,101.2,74.2,"1,032.8",246.8,452.3,25.5,991,55.8 +Eicher Motors Ltd.,EICHERMOT,505200,AUTOMOBILES & AUTO COMPONENTS,2/3 WHEELERS,"4,388.3","3,027.4","1,087.2",26.42%,142.5,12.7,"1,205.7",291.1,"1,016.2",37.1,"3,581",130.8 +Emami Ltd.,EMAMILTD,531162,FMCG,PERSONAL PRODUCTS,876,631.2,233.7,27.02%,46.1,2.2,196.4,15.8,178.5,4.1,697.8,16 +Endurance Technologies Ltd.,ENDURANCE,540153,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"2,560.5","2,226.7",318.3,12.51%,118.4,9.8,205.6,51.1,154.6,11,562.8,40 +Engineers India Ltd.,ENGINERSIN,532178,COMMERCIAL SERVICES & SUPPLIES,CONSULTING SERVICES,833.6,691.3,98.5,12.47%,8.3,0.4,133.6,32.2,127.5,2.3,472.7,8.4 +Escorts Kubota Ltd.,ESCORTS,500495,AUTOMOBILES & AUTO COMPONENTS,COMMERCIAL VEHICLES,"2,154.4","1,798.6",260.7,12.66%,40.8,3.1,311.9,79.7,223.3,20.6,910.5,82.4 +Exide Industries Ltd.,EXIDEIND,500086,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"4,408.9","3,872.4",499.1,11.42%,141.5,29.7,365.3,95.2,269.4,3.2,872.7,10.3 +Federal Bank Ltd.,FEDERALBNK,500469,BANKING AND FINANCE,BANKS,"6,548.2","1,603.8","1,400.3",24.18%,0,"3,544.1","1,342.7",342.6,994.1,4.3,"3,671.4",15.6 +Finolex Cables Ltd.,FINCABLES,500144,CONSUMER DURABLES,OTHER ELECTRICAL EQUIPMENT/PRODUCTS,"1,229.3","1,041.3",146.1,12.30%,10.8,0.4,176.7,52.3,154.2,10.1,643.9,42.1 +Finolex Industries Ltd.,FINPIPE,500940,GENERAL INDUSTRIALS,PLASTIC PRODUCTS,944.5,780.2,103,11.66%,27.4,12.5,124.5,35.4,98,1.6,459.3,7.4 +Firstsource Solutions Ltd.,FSL,532809,SOFTWARE & SERVICES,BPO/KPO,"1,556.9","1,311.2",228.8,14.86%,65.4,26.1,154.3,27.8,126.5,1.9,551.7,7.9 +GAIL (India) Ltd.,GAIL,532155,UTILITIES,UTILITIES,"33,191","29,405.5","3,580.2",10.85%,837.3,199.6,"2,748.7",696.3,"2,444.1",3.7,"5,283.8",8 +GlaxoSmithKline Pharmaceuticals Ltd.,GLAXO,500660,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,985.2,667.5,289.5,30.25%,18.1,0.4,299.2,81.7,217.5,12.8,647.8,38.2 +Glenmark Pharmaceuticals Ltd.,GLENMARK,532296,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"3,209.1","2,745.1",462.3,14.41%,141.5,121.5,-124.4,55.9,-81.9,-2.9,-196.3,-7 +Godrej Consumer Products Ltd.,GODREJCP,532424,FMCG,PERSONAL PRODUCTS,"3,667.9","2,897.8",704.2,19.55%,60.9,77.3,619.4,186.6,432.8,4.2,"1,750.1",17.1 +Godrej Industries Ltd.,GODREJIND,500164,DIVERSIFIED,DIVERSIFIED,"4,256.9","3,672.1",265.5,6.74%,89.3,333.1,162.4,75.9,87.3,2.6,880,26.1 +Godrej Properties Ltd.,GODREJPROP,533150,REALTY,REALTY,605.1,404.7,-61.7,-17.98%,7.4,48,145.1,38.8,66.8,2.4,662.6,23.8 +Granules India Ltd.,GRANULES,532482,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"1,191",976.5,213,17.90%,52.5,26,136,33.9,102.1,4.2,393.9,16.3 +Great Eastern Shipping Company Ltd.,GESHIP,500620,TRANSPORTATION,SHIPPING,"1,461.5",585.6,643.4,52.35%,186.7,77.1,611.9,17.3,594.7,41.6,"2,520.1",176.5 +Gujarat Alkalies & Chemicals Ltd.,GUJALKALI,530001,CHEMICALS & PETROCHEMICALS,COMMODITY CHEMICALS,"1,042.3",926.1,45.2,4.65%,95.2,10.8,10.2,-0.1,-18.4,-2.5,82.7,11.3 +Gujarat Gas Ltd.,GUJGASLTD,539336,UTILITIES,UTILITIES,"4,019.3","3,494.5",496.6,12.44%,117.9,7.8,399.1,102.9,296.2,4.3,"1,254.3",18.2 +Gujarat Narmada Valley Fertilizers & Chemicals Ltd.,GNFC,500670,FERTILIZERS,FERTILIZERS,"2,232","1,911",169,8.12%,78,1,242,64,182,11.7,932,60.1 +Gujarat Pipavav Port Ltd.,GPPL,533248,TRANSPORTATION,MARINE PORT & SERVICES,270.4,102,150.6,59.64%,28.8,2.2,141.1,53.4,92.3,1.9,341.8,7.1 +Gujarat State Fertilizer & Chemicals Ltd.,GSFC,500690,FERTILIZERS,FERTILIZERS,"3,313.2","2,881.4",237.3,7.61%,45.7,1.6,387,78.1,308.9,7.8,"1,056.2",26.5 +Gujarat State Petronet Ltd.,GSPL,532702,UTILITIES,UTILITIES,"4,455.9","3,497.2",913.7,20.72%,165,14.5,779.2,198.7,454.6,8.1,"1,522",27 +HCL Technologies Ltd.,HCLTECH,532281,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"27,037","20,743","5,929",22.23%,"1,01",156,"5,128","1,295","3,832",14.2,"15,445",56.9 +HDFC Bank Ltd.,HDFCBANK,500180,BANKING AND FINANCE,BANKS,"107,566.6","42,037.6","24,279.1",32.36%,0,"41,249.9","20,967.4","3,655","16,811.4",22.2,"54,474.6",71.8 +Havells India Ltd.,HAVELLS,517354,CONSUMER DURABLES,OTHER ELECTRICAL EQUIPMENT/PRODUCTS,"3,952.8","3,527",373.4,9.57%,81.2,9.3,335.3,86.2,249.1,4,"1,177.7",18.8 +Hero MotoCorp Ltd.,HEROMOTOCO,500182,AUTOMOBILES & AUTO COMPONENTS,2/3 WHEELERS,"9,741.2","8,173.5","1,359.5",14.26%,187.1,25,"1,355.6",353.1,"1,006.3",50.3,"3,247.6",162.5 +HFCL Ltd.,HFCL,500183,TELECOMMUNICATIONS EQUIPMENT,TELECOM CABLES,"1,128.7",978.9,132.6,11.93%,21.4,34.8,93.5,24,69.4,0.5,305.5,2.1 +Hindalco Industries Ltd.,HINDALCO,500440,METALS & MINING,ALUMINIUM AND ALUMINIUM PRODUCTS,"54,632","48,557","5,612",10.36%,"1,843","1,034","3,231","1,035","2,196",9.9,"8,423",37.9 +Hindustan Copper Ltd.,HINDCOPPER,513599,METALS & MINING,COPPER,392.6,260.2,121.2,31.77%,45.6,4.1,82.6,21.9,60.7,0.6,320.5,3.3 +Hindustan Petroleum Corporation Ltd.,HINDPETRO,500104,OIL & GAS,REFINERIES/PETRO-PRODUCTS,"96,093.4","87,512","8,24",8.61%,"1,247.3",590,"6,744.1","1,616","5,827",41.1,"16,645",117.3 +Hindustan Unilever Ltd.,HINDUNILVR,500696,FMCG,PERSONAL PRODUCTS,"15,806","11,826","3,797",24.30%,297,88,"3,59",931,"2,656",11.3,"10,284",43.8 +Hindustan Zinc Ltd.,HINDZINC,500188,METALS & MINING,ZINC,"7,014","3,652","3,139",46.22%,825,232,"2,305",576,"1,729",4.1,"8,432",20 +Housing and Urban Development Corporation Ltd.,HUDCO,540530,BANKING AND FINANCE,HOUSING FINANCE,"1,880.8",82.7,"1,809.6",97.04%,2.4,"1,216.8",606.4,154.7,451.6,2.3,"1,790.7",8.9 +ITC Ltd.,ITC,500875,FOOD BEVERAGES & TOBACCO,CIGARETTES-TOBACCO PRODUCTS,"18,439.3","11,320.2","6,454.2",36.31%,453,9.9,"6,656.2","1,700.3","4,898.1",3.9,"20,185.1",16.2 +ICICI Bank Ltd.,ICICIBANK,532174,BANKING AND FINANCE,BANKS,"57,292.3","23,911","15,473.2",39.74%,0,"17,908","14,824.2","3,808.8","11,805.6",15.6,"41,086.8",58.7 +ICICI Prudential Life Insurance Company Ltd.,ICICIPRULI,540133,BANKING AND FINANCE,LIFE INSURANCE,"17,958.1","17,612.3",-229.6,-1.32%,0,0,340.2,32.5,243.9,1.7,906.9,6.3 +IDBI Bank Ltd.,IDBI,500116,BANKING AND FINANCE,BANKS,"7,063.7","1,922.3","2,175.3",36.02%,0,"2,966.1","2,396.9","1,003.7","1,385.4",1.3,"4,776.3",4.4 +IDFC First Bank Ltd.,IDFCFIRSTB,539437,BANKING AND FINANCE,BANKS,"8,765.8","3,849","1,511.2",20.54%,0,"3,405.6",982.8,236,746.9,1.1,"2,911.1",4.3 +IDFC Ltd.,IDFC,532659,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),36.7,6,30.6,83.56%,0,0,30.6,6.6,223.5,1.4,"4,147.1",25.9 +IRB Infrastructure Developers Ltd.,IRB,532947,CEMENT AND CONSTRUCTION,ROADS & HIGHWAYS,"1,874.5",950.4,794.6,45.54%,232.7,434.6,256.9,85.8,95.7,0.2,501,0.8 +ITI Ltd.,ITI,523610,TELECOMMUNICATIONS EQUIPMENT,TELECOM EQUIPMENT,256.1,299.3,-52.8,-21.42%,13.3,69.3,-125.8,0,-126,-1.3,-388.4,-4 +Vodafone Idea Ltd.,IDEA,532822,TELECOM SERVICES,TELECOM SERVICES,"10,750.8","6,433.5","4,282.8",39.97%,"5,667.3","6,569","-7,919",817.7,"-8,737.9",-1.8,"-30,986.8",-6.4 +India Cements Ltd.,INDIACEM,530005,CEMENT AND CONSTRUCTION,CEMENT & CEMENT PRODUCTS,"1,272.4","1,26",4.4,0.35%,55,60.4,-103,-17.4,-80.1,-2.6,-261.1,-8.4 +Indiabulls Housing Finance Ltd.,IBULHSGFIN,535789,BANKING AND FINANCE,HOUSING FINANCE,"2,242.3",190.6,"1,779.2",79.88%,22.9,"1,349.8",421.6,123.6,298,6.5,"1,146",24.3 +Indian Bank,INDIANB,532814,BANKING AND FINANCE,BANKS,"15,929.4","3,599.1","4,327.7",31.44%,0,"8,002.6","2,776.7",768.6,"2,068.5",16.6,"6,893.3",55.3 +Indian Hotels Company Ltd.,INDHOTEL,500850,HOTELS RESTAURANTS & TOURISM,HOTELS,"1,480.9","1,078.4",354.8,24.75%,111.2,59,232.2,72.3,166.9,1.2,"1,100.3",7.7 +Indian Oil Corporation Ltd.,IOC,530965,OIL & GAS,OIL MARKETING & DISTRIBUTION,"179,752.1","156,013.1","23,328.4",13.01%,"3,609.6","2,135","18,090.2","4,699.7","13,114.3",9.5,"38,614.3",27.3 +Indian Overseas Bank,IOB,532388,BANKING AND FINANCE,BANKS,"6,941.5","1,785.1","1,679.8",28.84%,0,"3,476.6",635.5,8.3,627.2,0.3,"2,341.9",1.2 +Indraprastha Gas Ltd.,IGL,532514,UTILITIES,UTILITIES,"3,520.2","2,801.6",656.9,18.99%,102.2,2.5,613.9,151.4,552.7,7.9,"1,806.2",25.8 +IndusInd Bank Ltd.,INDUSINDBK,532187,BANKING AND FINANCE,BANKS,"13,529.7","3,449.9","3,908.7",34.75%,0,"6,171.1","2,934.9",732.9,"2,202.2",28.4,"8,333.7",107.2 +Info Edge (India) Ltd.,NAUKRI,532777,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,792,421.2,204.7,32.70%,25.9,8.2,382.8,68.7,205.1,15.9,-25.6,-2 +InterGlobe Aviation Ltd.,INDIGO,539448,TRANSPORTATION,AIRLINES,"15,502.9","12,743.6","2,200.3",14.72%,"1,549","1,021.3",189.1,0.2,188.9,4.9,"5,621.3",145.7 +Ipca Laboratories Ltd.,IPCALAB,524494,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"2,072.5","1,712.7",321.3,15.80%,90.3,44.1,225.4,87.9,145.1,5.7,492.2,19.4 +J B Chemicals & Pharmaceuticals Ltd.,JBCHEPHARM,506943,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,889.4,638.2,243.5,27.62%,32.2,10.4,208.7,58.1,150.6,9.7,486.6,31.4 +JK Cement Ltd.,JKCEMENT,532644,CEMENT AND CONSTRUCTION,CEMENT & CEMENT PRODUCTS,"2,782.1","2,285.8",467,16.96%,137.1,115,244.2,65.7,178.1,23.1,444,57.5 +JK Lakshmi Cement Ltd.,JKLAKSHMI,500380,CEMENT AND CONSTRUCTION,CEMENT & CEMENT PRODUCTS,"1,588.5","1,357.3",217.3,13.80%,56.6,33.6,141,45.1,92.7,7.9,357.6,30.4 +JM Financial Ltd.,JMFINANCIL,523405,DIVERSIFIED,HOLDING COMPANIES,"1,214",407.9,662.6,55.34%,13.2,388.1,277.9,72.4,194.9,2,608.1,6.4 +JSW Energy Ltd.,JSWENERGY,533148,UTILITIES,ELECTRIC UTILITIES,"3,387.4","1,379","1,880.4",57.69%,408.7,513.7,"1,085.9",235.1,850.2,5.2,"1,591.7",9.7 +JSW Steel Ltd.,JSWSTEEL,500228,METALS & MINING,IRON & STEEL/INTERM.PRODUCTS,"44,821","36,698","7,886",17.69%,"2,019","2,084","4,609","1,812","2,76",11.4,"9,252",38.1 +Jindal Stainless Ltd.,JSL,532508,METALS & MINING,IRON & STEEL/INTERM.PRODUCTS,"9,829","8,566.5","1,230.6",12.56%,221.9,155.6,985.7,229.1,774.3,9.4,"2,600.2",31.6 +Jindal Steel & Power Ltd.,JINDALSTEL,532286,METALS & MINING,IRON & STEEL/INTERM.PRODUCTS,"12,282","9,964.5","2,285.7",18.66%,603.7,329.4,"1,384.5",-5.8,"1,387.8",13.8,"4,056",40.4 +Jubilant Foodworks Ltd.,JUBLFOOD,533155,HOTELS RESTAURANTS & TOURISM,RESTAURANTS,"1,375.7","1,091.4",277.2,20.25%,141.9,56.8,85.5,23.3,97.2,1.5,235,3.6 +Just Dial Ltd.,JUSTDIAL,535648,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,318.5,211.8,48.8,18.71%,12.2,2.4,92.1,20.3,71.8,8.4,314.1,36.9 +Jyothy Labs Ltd.,JYOTHYLAB,532926,FMCG,PERSONAL PRODUCTS,745.6,597,135.4,18.48%,12.3,1.2,135.1,31.1,104.2,2.8,326.9,8.9 +KRBL Ltd.,KRBL,530813,FMCG,PACKAGED FOODS,"1,246.5","1,018.9",194.5,16.03%,19.9,0.8,206.8,53.6,153.3,6.5,671.4,29.3 +Kajaria Ceramics Ltd.,KAJARIACER,500233,DIVERSIFIED CONSUMER SERVICES,FURNITURE-FURNISHING-PAINTS,"1,129.9",941.9,179.7,16.02%,36.1,4.3,147.7,36.6,108,6.8,397.8,25 +Kalpataru Projects International Ltd.,KPIL,522287,UTILITIES,ELECTRIC UTILITIES,"4,53","4,148",370,8.19%,113,137,132,42,89,5.5,478,29.9 +Kansai Nerolac Paints Ltd.,KANSAINER,500165,DIVERSIFIED CONSUMER SERVICES,FURNITURE-FURNISHING-PAINTS,"1,978.6","1,683.3",273.2,13.97%,47.4,7.6,240.3,64.8,177.2,2.2,"1,118.8",13.8 +Karur Vysya Bank Ltd.,KARURVYSYA,590003,BANKING AND FINANCE,BANKS,"2,336",616.4,637.9,31.94%,0,"1,081.7",511.5,133.1,378.4,4.7,"1,364.2",17 +KEC International Ltd.,KEC,532714,GENERAL INDUSTRIALS,HEAVY ELECTRICAL EQUIPMENT,"4,514.9","4,224.7",274.3,6.10%,46.5,177.8,65.8,9.9,55.8,2.2,187.9,7.3 +Kotak Mahindra Bank Ltd.,KOTAKBANK,500247,BANKING AND FINANCE,BANKS,"21,559.5","9,681","6,343",46.24%,0,"5,535.5","5,888.3","1,465.5","4,461",22.4,"17,172.7",86.4 +L&T Finance Holdings Ltd.,L&TFH,533519,DIVERSIFIED,HOLDING COMPANIES,"3,482.1",935.3,"1,882.4",58.57%,28.3,"1,324.9",797.4,203.2,595.1,2.4,"2,080.8",8.4 +L&T Technology Services Ltd.,LTTS,540115,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"2,427.7","1,910.9",475.6,19.93%,68.1,12.6,436.1,120.2,315.4,29.8,"1,239.7",117.5 +LIC Housing Finance Ltd.,LICHSGFIN,500253,BANKING AND FINANCE,HOUSING FINANCE,"6,765.9",250.6,"6,095.7",90.10%,13.2,"4,599.9","1,483",291.2,"1,193.5",21.7,"4,164.5",75.7 +Lakshmi Machine Works Ltd.,LAXMIMACH,500252,GENERAL INDUSTRIALS,INDUSTRIAL MACHINERY,"1,355.5","1,184.5",136,10.30%,23.6,0,147.4,32.3,115.1,107.8,416,389.5 +Laurus Labs Ltd.,LAURUSLABS,540222,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"1,226.2","1,036.6",187.9,15.34%,93.4,42.4,53.9,14.6,37,0.7,367.8,6.8 +Lupin Ltd.,LUPIN,500257,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"5,079","4,120.8",917.8,18.21%,247.8,80.6,629.7,134.3,489.5,10.8,"1,331.2",29.2 +MMTC Ltd.,MMTC,513377,COMMERCIAL SERVICES & SUPPLIES,COMMODITY TRADING & DISTRIBUTION,-167.2,-180.1,-30.4,14.42%,0.8,1.1,12.1,1.5,52,0.3,174.1,1.2 +MRF Ltd.,MRF,500290,AUTOMOBILES & AUTO COMPONENTS,AUTO TYRES & RUBBER PRODUCTS,"6,287.8","5,060.2","1,156.9",18.61%,351.5,85.5,790.6,203.9,586.7,1383.3,"1,690.9",3988 +Mahanagar Gas Ltd.,MGL,539957,UTILITIES,UTILITIES,"1,772.7","1,250.1",478.9,27.70%,65.8,2.5,454.3,115.8,338.5,34.3,"1,147.8",116.2 +Mahindra & Mahindra Financial Services Ltd.,M&MFIN,532720,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),"3,863.5","1,077.5","2,109.3",55.03%,67.1,"1,703.4",369.1,96,281.1,2.3,"1,982.5",16 +Mahindra & Mahindra Ltd.,M&M,500520,AUTOMOBILES & AUTO COMPONENTS,CARS & UTILITY VEHICLES,"35,027.2","28,705.9","5,729.6",16.64%,"1,138.6","1,835.2","3,347.5","1,083.7","2,347.8",21.1,"11,169.4",100.2 +Mahindra Holidays & Resorts India Ltd.,MHRIL,533088,HOTELS RESTAURANTS & TOURISM,HOTELS,672.2,519.3,136,20.76%,83.8,33.3,35.8,14,21.3,1.1,66,3.3 +Manappuram Finance Ltd.,MANAPPURAM,531213,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),"2,174",555.6,"1,481.3",68.68%,62.5,689.4,746.7,186.1,558.4,6.6,"1,859.8",22 +Mangalore Refinery And Petrochemicals Ltd.,MRPL,500109,OIL & GAS,REFINERIES/PETRO-PRODUCTS,"22,904.7","20,705.6","2,138.2",9.36%,296,311.2,"1,592",546.2,"1,051.7",6,"3,784.9",21.6 +Marico Ltd.,MARICO,531642,FMCG,PERSONAL PRODUCTS,"2,514","1,979",497,20.07%,39,20,476,116,353,2.7,"1,41",10.9 +Maruti Suzuki India Ltd.,MARUTI,532500,AUTOMOBILES & AUTO COMPONENTS,CARS & UTILITY VEHICLES,"37,902.1","32,282.5","4,790.3",12.92%,794.4,35.1,"4,790.1","1,083.8","3,764.3",124.6,"11,351.8",375.9 +Max Financial Services Ltd.,MFSL,500271,BANKING AND FINANCE,LIFE INSURANCE,"10,189.1","10,024.6",143.9,1.42%,0.8,9.4,158.2,-12.1,147.9,4.3,506.4,14.7 +UNO Minda Ltd.,UNOMINDA,532539,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"3,630.2","3,219.8",401.6,11.09%,125.4,27.2,257.9,73.3,225,3.9,742.4,13 +Motilal Oswal Financial Services Ltd.,MOTILALOFS,532892,BANKING AND FINANCE,OTHER FINANCIAL SERVICES,"1,650.7",724.1,904.5,55.18%,17.3,241.1,657.6,124.2,531.2,35.9,"1,449.3",97.8 +MphasiS Ltd.,MPHASIS,526299,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"3,325.5","2,680.9",595.6,18.18%,89,34,521.7,129.7,391.9,20.8,"1,605.6",85.1 +Muthoot Finance Ltd.,MUTHOOTFIN,533398,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),"3,631.9",723.4,"2,801.6",77.69%,22.2,"1,335","1,470.2",374.9,"1,059.6",26.4,"3,982.9",99.2 +Natco Pharma Ltd.,NATCOPHARM,524816,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"1,060.8",573.4,458,44.41%,43.6,4.2,439.6,70.6,369,20.6,"1,127.4",63 +NBCC (India) Ltd.,NBCC,534309,CEMENT AND CONSTRUCTION,CONSTRUCTION & ENGINEERING,"2,129.1","1,957.7",95.5,4.65%,1.3,0,104.6,22.9,79.6,0.4,332.2,1.8 +NCC Ltd.,NCC,500294,CEMENT AND CONSTRUCTION,CONSTRUCTION & ENGINEERING,"4,746.4","4,415.9",303.7,6.44%,53.2,153.5,123.8,38.8,77.3,1.2,599.4,9.5 +NHPC Ltd.,NHPC,533098,UTILITIES,ELECTRIC UTILITIES,"3,113.8","1,173.9","1,757.4",59.95%,294.9,104.8,"1,618.3",-75,"1,545.8",1.5,"3,897.8",3.9 +Coforge Ltd.,COFORGE,532541,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"2,285.1","1,935.3",340.9,14.98%,77.2,31.9,240.7,52.8,187.9,29.6,696.2,113.2 +NLC India Ltd.,NLCINDIA,513683,UTILITIES,ELECTRIC UTILITIES,"3,234","2,143",834.6,28.03%,455.1,213.9,"1,700.6",614.7,"1,084.7",7.8,"1,912.3",13.8 +NTPC Ltd.,NTPC,532555,UTILITIES,ELECTRIC UTILITIES,"45,384.6","32,303.2","12,680.2",28.19%,"4,037.7","2,920.5","6,342.9","2,019.7","4,614.6",4.8,"19,125.2",19.7 +Narayana Hrudayalaya Ltd.,NH,539551,DIVERSIFIED CONSUMER SERVICES,HEALTHCARE FACILITIES,"1,323.6",997.1,308.1,23.61%,55.3,22.9,248.4,21.7,226.6,11.2,737.5,36.1 +National Aluminium Company Ltd.,NATIONALUM,532234,METALS & MINING,ALUMINIUM AND ALUMINIUM PRODUCTS,"3,112","2,646.9",396.5,13.03%,186.2,4,275,68.7,187.3,1,"1,272.4",6.9 +Navin Fluorine International Ltd.,NAVINFLUOR,532504,CHEMICALS & PETROCHEMICALS,COMMODITY CHEMICALS,494.9,373.4,98.3,20.84%,24.2,20,77.2,16.6,60.6,12.2,365,73.7 +Oberoi Realty Ltd.,OBEROIRLTY,533273,REALTY,REALTY,"1,243.8",579.2,638.2,52.42%,11.3,56.5,596.8,142.1,456.8,12.6,"1,961.3",53.9 +Oil And Natural Gas Corporation Ltd.,ONGC,500312,OIL & GAS,EXPLORATION & PRODUCTION,"149,388.5","118,618.4","28,255.3",19.24%,"6,698.1","2,603.3","21,564.9","5,633.6","13,734.1",10.9,"43,072.5",34.2 +Oil India Ltd.,OIL,533106,OIL & GAS,EXPLORATION & PRODUCTION,"9,200.1","5,293.3","3,523.2",39.96%,499,278.9,762,67.6,420.7,3.9,"5,874.5",54.2 +Oracle Financial Services Software Ltd.,OFSS,532466,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"1,509.6",886.4,558.1,38.64%,19,8,596.2,178.8,417.4,48.2,"1,835.1",211.9 +PI Industries Ltd.,PIIND,523642,CHEMICALS & PETROCHEMICALS,AGROCHEMICALS,"2,163.8","1,565.5",551.4,26.05%,80.3,7.8,510.2,31.7,480.5,31.7,"1,495.8",98.4 +PNB Housing Finance Ltd.,PNBHOUSING,540173,BANKING AND FINANCE,HOUSING FINANCE,"1,779.4",158.8,"1,574.1",88.54%,11.3,"1,057.3",507.1,124.1,383,14.8,"1,278.7",49.3 +PNC Infratech Ltd.,PNCINFRA,539150,CEMENT AND CONSTRUCTION,ROADS & HIGHWAYS,"1,932.4","1,511.6",399.8,20.92%,40.9,161.3,218.6,70.7,147.9,5.8,614.3,23.9 +PVR INOX Ltd.,PVRINOX,532689,RETAILING,SPECIALTY RETAIL,"2,023.7","1,293.1",706.8,35.34%,308.6,200.3,221.7,55.5,166.3,17,-232.5,-23.7 +Page Industries Ltd.,PAGEIND,532827,TEXTILES APPARELS & ACCESSORIES,OTHER APPARELS & ACCESSORIES,"1,126.8",891.6,233.5,20.76%,24.6,11.2,199.4,49.1,150.3,134.7,510.7,457.9 +Persistent Systems Ltd.,PERSISTENT,533179,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"2,449","2,006.5",405.2,16.80%,74.4,12.3,355.8,92.5,263.3,35,981.5,127.6 +Petronet LNG Ltd.,PETRONET,532522,OIL & GAS,OIL MARKETING & DISTRIBUTION,"12,686.2","11,317.9","1,214.7",9.69%,194.8,74.7,"1,098.8",283.9,855.7,5.7,"3,490.3",23.3 +Pfizer Ltd.,PFIZER,500680,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,611.3,392.6,182.6,31.75%,15.4,2.7,200.5,51.6,149,32.6,522.8,114.3 +Phoenix Mills Ltd.,PHOENIXLTD,503100,REALTY,REALTY,906.6,361.2,506,57.82%,65.9,96.5,375.2,71.4,252.6,14.2,923.6,51.7 +Pidilite Industries Ltd.,PIDILITIND,500331,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,"3,107.6","2,396.3",679.7,22.10%,75.2,13.1,623,163.1,450.1,8.8,"1,505.5",29.6 +Power Finance Corporation Ltd.,PFC,532810,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),"22,403.7",315.4,"22,941.9",102.46%,12.7,"14,313.1","8,628.8","2,000.6","4,833.1",14.7,"17,946.4",54.4 +Power Grid Corporation of India Ltd.,POWERGRID,532898,UTILITIES,ELECTRIC UTILITIES,"11,530.4","1,358.7","9,908.4",87.94%,"3,277","2,341.3","4,393.4",573.7,"3,781.4",4.1,"15,344.4",16.5 +Prestige Estates Projects Ltd.,PRESTIGE,ASM,REALTY,REALTY,"3,256","1,643.9",592.5,26.49%,174.1,263.9,"1,174.1",256.4,850.9,21.2,"1,714",42.8 +Prism Johnson Ltd.,PRSMJOHNSN,500338,CEMENT AND CONSTRUCTION,CEMENT & CEMENT PRODUCTS,"1,846","1,745.4",92.4,5.03%,95.2,43.5,210,30.4,182.7,3.6,154.2,3.1 +Procter & Gamble Hygiene & Healthcare Ltd.,PGHH,500459,FMCG,PERSONAL PRODUCTS,"1,154.1",853.5,284.9,25.03%,14.3,1.9,284.5,73.8,210.7,64.9,734.4,226.3 +Punjab National Bank,PNB,532461,BANKING AND FINANCE,BANKS,"29,857","6,798.1","6,239.1",23.23%,0,"16,819.8","2,778.3","1,013.8","1,990.2",1.8,"5,904.8",5.4 +Quess Corp Ltd.,QUESS,539978,SOFTWARE & SERVICES,BPO/KPO,"4,763.5","4,584.8",163.6,3.44%,69.7,28.1,79.3,8.3,71.9,4.8,240.9,16.2 +RBL Bank Ltd.,RBLBANK,540065,BANKING AND FINANCE,BANKS,"3,720.6","1,422.6",765.4,25.45%,0,"1,532.6",125,-206.1,331.1,5.5,"1,173.9",19.5 +Radico Khaitan Ltd.,RADICO,532497,FOOD BEVERAGES & TOBACCO,BREWERIES & DISTILLERIES,925.7,803.8,121.2,13.10%,26.1,12.5,83.3,21.4,64.8,4.8,237,17.7 +Rain Industries Ltd.,RAIN,500339,CHEMICALS & PETROCHEMICALS,PETROCHEMICALS,"4,208.9","3,794.3",366,8.80%,192.5,241.7,-19.5,46.2,-90.2,-2.7,270.4,8 +Rajesh Exports Ltd.,RAJESHEXPO,531500,TEXTILES APPARELS & ACCESSORIES,GEMS & JEWELLERY,"38,079.4","38,015.8",50.1,0.13%,10.7,0,53,7.7,45.3,1.5,"1,142.2",38.7 +Rallis India Ltd.,RALLIS,500355,CHEMICALS & PETROCHEMICALS,AGROCHEMICALS,837,699,133,15.99%,26,3,110,28,82,4.2,98.4,5.2 +Rashtriya Chemicals & Fertilizers Ltd.,RCF,524230,FERTILIZERS,FERTILIZERS,"4,222.1","4,049.3",105.9,2.55%,56.1,44,72.8,21.1,51,0.9,523.6,9.5 +Redington Ltd.,REDINGTON,532805,COMMERCIAL SERVICES & SUPPLIES,COMMODITY TRADING & DISTRIBUTION,"22,296.6","21,738.7",481.4,2.17%,43.7,105.8,408.3,96.7,303.5,3.9,"1,242",15.9 +Relaxo Footwears Ltd.,RELAXO,530517,RETAILING,FOOTWEAR,725.9,623.8,91.5,12.79%,36.9,4.7,60.4,16.2,44.2,1.8,193.9,7.8 +Reliance Industries Ltd.,RELIANCE,500325,OIL & GAS,REFINERIES/PETRO-PRODUCTS,"238,797","193,988","40,968",17.44%,"12,585","5,731","26,493","6,673","17,394",25.7,"68,496",101.2 +REC Ltd.,RECLTD,532955,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),"11,701.3",275.1,"12,180.5",104.21%,6.1,"7,349.8","4,837.6","1,047.7","3,789.9",14.4,"12,738.6",48.4 +SJVN Ltd.,SJVN,533206,UTILITIES,ELECTRIC UTILITIES,951.6,172.2,706.2,80.40%,101.9,124.2,567.7,129.2,439.6,1.1,"1,016",2.6 +SKF India Ltd.,SKFINDIA,500472,GENERAL INDUSTRIALS,OTHER INDUSTRIAL GOODS,"1,145.5","1,003.7",121.5,10.80%,19.3,0.5,122,31.7,90,18.2,484,97.9 +SRF Ltd.,SRF,503806,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,"3,206.5","2,551.2",626.2,19.71%,161.2,79.3,414.8,114,300.8,10.2,"1,733.4",58.5 +Sanofi India Ltd.,SANOFI,500674,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,726.4,506.1,208.5,29.17%,9.9,0.3,210.1,57.9,152.1,66.1,596.3,259.3 +Schaeffler India Ltd.,SCHAEFFLER,505790,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"1,879.2","1,506.3",342,18.50%,55.6,1.6,315.7,80.7,235,15,922.6,59 +Shree Cements Ltd.,SHREECEM,500387,CEMENT AND CONSTRUCTION,CEMENT & CEMENT PRODUCTS,"4,932.1","3,914.1",886,18.46%,411.7,67,539.2,92.6,446.6,123.8,"1,826.8",506.3 +Shriram Finance Ltd.,SHRIRAMFIN,511218,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),"8,893","1,409.4","6,334.3",71.30%,141.4,"3,798","2,404.2",614.9,"1,786.1",47.6,"6,575.4",175.2 +Siemens Ltd.,SIEMENS,500550,GENERAL INDUSTRIALS,HEAVY ELECTRICAL EQUIPMENT,"5,953.2","5,107.5",700.2,12.06%,78.6,4.9,762.2,190.5,571.3,16.1,"1,960.9",55.1 +Sobha Ltd.,SOBHA,532784,REALTY,REALTY,773.6,665.8,75.4,10.18%,19.3,63.9,24.7,9.7,14.9,1.6,107.4,11.3 +Solar Industries India Ltd.,SOLARINDS,532725,GENERAL INDUSTRIALS,OTHER INDUSTRIAL PRODUCTS,"1,355.2","1,011.3",336.1,24.95%,33.7,24.9,285.3,75.5,200.1,22.1,808.2,89.3 +Sonata Software Ltd.,SONATSOFTW,532221,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"1,935.8","1,715.2",197.3,10.32%,33.3,20.7,166.5,42.3,124.2,9,475.7,34.3 +State Bank of India,SBIN,500112,BANKING AND FINANCE,BANKS,"144,256.1","58,597.6","22,703.3",21.14%,0,"62,955.2","21,935.7","5,552.5","17,196.2",18,"69,304.1",77.7 +Steel Authority of India (SAIL) Ltd.,SAIL,500113,METALS & MINING,IRON & STEEL/INTERM.PRODUCTS,"29,858.2","25,836.7","3,875.4",13.04%,"1,326.6",605.2,"1,674.7",464.2,"1,305.6",3.2,"3,219.5",7.8 +Sun Pharma Advanced Research Company Ltd.,SPARC,532872,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,29.7,112.7,-91.5,-431.87%,3.2,0.3,-86.4,0,-86.4,-2.7,-253.6,-7.8 +Sun Pharmaceutical Industries Ltd.,SUNPHARMA,524715,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"12,486","9,013","3,179.4",26.08%,632.8,49.3,"2,790.9",390.1,"2,375.5",9.9,"8,548.5",35.6 +Sun TV Network Ltd.,SUNTV,532733,MEDIA,BROADCASTING & CABLE TV,"1,160.2",320.6,727.8,69.42%,218.8,1.7,619.1,154.4,464.7,11.8,"1,861.8",47.2 +Sundram Fasteners Ltd.,SUNDRMFAST,500403,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"1,429.1","1,191.1",230.7,16.23%,54.5,7.4,176.2,43.1,131.9,6.3,502.9,23.9 +Sunteck Realty Ltd.,SUNTECK,512179,REALTY,REALTY,36.2,39.1,-14.1,-56.70%,2.2,15.8,-20.9,-6.4,-13.9,-1,-46.5,-3.3 +Supreme Industries Ltd.,SUPREMEIND,509930,GENERAL INDUSTRIALS,PLASTIC PRODUCTS,"2,321.4","1,952.5",356.2,15.43%,71.9,1.6,295.4,76.3,243.2,19.1,"1,028.2",80.9 +Suzlon Energy Ltd.,SUZLON,ASM,GENERAL INDUSTRIALS,HEAVY ELECTRICAL EQUIPMENT,"1,428.7","1,196.4",225,15.83%,51.2,43.7,102.4,0.1,102.3,0.1,561.4,0.4 +Syngene International Ltd.,SYNGENE,539268,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,931.7,656,254.1,27.92%,104.6,13,150.7,34.2,116.5,2.9,498.3,12.4 +TTK Prestige Ltd.,TTKPRESTIG,517506,CONSUMER DURABLES,HOUSEWARE,747.2,648.6,80.8,11.08%,15.9,3.1,79.5,20.5,59.3,4.3,224.3,16.2 +TV18 Broadcast Ltd.,TV18BRDCST,532800,MEDIA,BROADCASTING & CABLE TV,"1,989","1,992.2",-198.1,-11.04%,50.1,33.8,-87.1,-6.5,-28.9,-0.2,92.2,0.5 +TVS Motor Company Ltd.,TVSMOTOR,532343,AUTOMOBILES & AUTO COMPONENTS,2/3 WHEELERS,"9,983.8","8,576.9","1,355.9",13.65%,237.1,483.3,686.4,259.8,386.3,8.1,"1,457.6",30.7 +Tata Consultancy Services Ltd.,TCS,532540,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"60,698","43,946","15,746",26.38%,"1,263",159,"15,33","3,95","11,342",31,"44,654",122 +Tata Elxsi Ltd.,TATAELXSI,500408,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,912.8,618.2,263.5,29.89%,25,5.8,263.9,63.8,200,32.1,785.1,126.1 +Tata Consumer Products Ltd.,TATACONSUM,500800,FMCG,PACKAGED FOODS,"3,823.6","3,196.7",537.1,14.38%,93.9,27.6,490.9,131.7,338.2,3.6,"1,275.2",13.7 +Tata Motors Limited (DVR),TATAMTRDVR,570001,AUTOMOBILES & AUTO COMPONENTS,COMMERCIAL VEHICLES,,,,,,,,,,,, +Tata Motors Ltd.,TATAMOTORS,500570,AUTOMOBILES & AUTO COMPONENTS,COMMERCIAL VEHICLES,"106,759","91,361.3","13,766.9",13.10%,"6,636.4","2,651.7","5,985.9","2,202.8","3,764",9.8,"15,332.3",40 +Tata Power Company Ltd.,TATAPOWER,500400,UTILITIES,ELECTRIC UTILITIES,"16,029.5","12,647","3,091",19.64%,925.9,"1,181.8",979.2,213.3,875.5,2.7,"3,570.8",11.2 +Tata Steel Ltd.,TATASTEEL,500470,METALS & MINING,IRON & STEEL/INTERM.PRODUCTS,"55,910.2","51,414.1","4,267.8",7.66%,"2,479.8","1,959.4","-6,842.1",-228,"-6,196.2",-5.1,"-6,081.3",-5 +Tech Mahindra Ltd.,TECHM,532755,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"13,128.1","11,941.1",922.8,7.17%,465.7,97.5,623.8,110,493.9,5.6,"3,600.7",40.9 +The Ramco Cements Ltd.,RAMCOCEM,500260,CEMENT AND CONSTRUCTION,CEMENT & CEMENT PRODUCTS,"2,352.1","1,935",405.6,17.33%,162.8,116.5,137.8,37,72,3.1,348.9,14.8 +Thermax Ltd.,THERMAX,500411,GENERAL INDUSTRIALS,HEAVY ELECTRICAL EQUIPMENT,"2,368.3","2,097.8",204.6,8.89%,33,19.8,217.7,58.9,157.7,14,498.8,44.3 +Timken India Ltd.,TIMKEN,522113,GENERAL INDUSTRIALS,OTHER INDUSTRIAL PRODUCTS,692.1,546.5,135.5,19.87%,21.1,0.9,123.6,30.6,93,12.4,358.3,47.6 +Titan Company Ltd.,TITAN,500114,TEXTILES APPARELS & ACCESSORIES,GEMS & JEWELLERY,"12,653","11,118","1,411",11.26%,144,140,"1,251",336,915,10.3,"3,302",37.1 +Torrent Pharmaceuticals Ltd.,TORNTPHARM,500420,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"2,686","1,835",825,31.02%,201,91,559,173,386,11.4,"1,334",39.4 +Torrent Power Ltd.,TORNTPOWER,532779,UTILITIES,ELECTRIC UTILITIES,"7,069.1","5,739.5","1,221.4",17.55%,341.7,247.2,740.7,198.1,525.9,10.9,"2,176.8",45.3 +Trent Ltd.,TRENT,500251,RETAILING,DEPARTMENT STORES,"3,062.5","2,525.8",456.6,15.31%,152.2,95.5,288.9,86.3,234.7,6.6,629.4,17.7 +Trident Ltd.,TRIDENT,521064,TEXTILES APPARELS & ACCESSORIES,TEXTILES,"1,812","1,557.3",240.3,13.37%,89.4,35,130.4,40.1,90.7,0.2,458.1,0.9 +UPL Ltd.,UPL,512070,CHEMICALS & PETROCHEMICALS,AGROCHEMICALS,"10,275","8,807","1,325",13.03%,657,871,-185,-96,-189,-2.5,"1,856",24.7 +UltraTech Cement Ltd.,ULTRACEMCO,532538,CEMENT AND CONSTRUCTION,CEMENT & CEMENT PRODUCTS,"16,179.3","13,461.2","2,550.9",15.93%,797.8,233.9,"1,686.2",409.4,"1,281.5",44.5,"5,694.1",197.2 +Union Bank of India,UNIONBANK,532477,BANKING AND FINANCE,BANKS,"28,952.5","6,189.3","7,265",29.38%,0,"15,498.2","5,492.3","1,944","3,571.8",5.1,"11,918.9",16.1 +United Breweries Ltd.,UBL,532478,FOOD BEVERAGES & TOBACCO,BREWERIES & DISTILLERIES,"1,902.1","1,705.8",184.3,9.75%,50.9,1.4,144,36.9,107.3,4.1,251.3,9.5 +United Spirits Ltd.,MCDOWELL-N,532432,FOOD BEVERAGES & TOBACCO,BREWERIES & DISTILLERIES,"6,776.6","6,269.8",466.7,6.93%,65.3,26.2,446,106.3,339.3,4.8,"1,133",15.6 +V-Guard Industries Ltd.,VGUARD,532953,CONSUMER DURABLES,OTHER ELECTRICAL EQUIPMENT/PRODUCTS,"1,147.9","1,041.3",92.5,8.16%,19.8,9.3,77.5,18.6,59,1.4,215.2,5 +Vardhman Textiles Ltd.,VTL,502986,TEXTILES APPARELS & ACCESSORIES,TEXTILES,"2,487","2,192.1",205.4,8.57%,103.7,22,169.2,41.7,134.3,4.7,531.9,18.7 +Varun Beverages Ltd.,VBL,540180,FOOD BEVERAGES & TOBACCO,NON-ALCOHOLIC BEVERAGES,"3,889","2,988.4",882.1,22.79%,170.8,62.5,667.3,152.9,501.1,3.9,"1,998.7",15.4 +Vinati Organics Ltd.,VINATIORGA,524200,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,464.4,337.3,110.8,24.73%,13.7,0.3,113,28.9,84.2,8.2,408.2,39.7 +Voltas Ltd.,VOLTAS,500575,CONSUMER DURABLES,CONSUMER ELECTRONICS,"2,363.7","2,222.5",70.3,3.06%,11.7,11.4,118.1,49.3,36.7,1.1,199.5,6 +ZF Commercial Vehicle Control Systems India Ltd.,ZFCVINDIA,533023,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"1,015.8",846.2,145.5,14.67%,27.1,1.3,141.2,35.5,105.7,55.7,392,206.7 +Welspun Corp Ltd.,WELCORP,ASM,METALS & MINING,IRON & STEEL PRODUCTS,"4,161.4","3,659.9",399.5,9.84%,85.7,75,340.8,79,384.7,14.7,809.2,30.9 +Welspun Living Ltd.,WELSPUNLIV,514162,TEXTILES APPARELS & ACCESSORIES,TEXTILES,"2,542.4","2,151.1",358,14.27%,98.5,33.8,258.9,58.7,196.7,2,526.1,5.4 +Whirlpool of India Ltd.,WHIRLPOOL,500238,CONSUMER DURABLES,CONSUMER ELECTRONICS,"1,555.5","1,448.4",73.2,4.81%,49.2,5.6,52.3,14.1,36.6,2.9,198.8,15.7 +Wipro Ltd.,WIPRO,507685,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"23,255.7","18,543.2","3,972.7",17.64%,897,303.3,"3,512.2",841.9,"2,646.3",5.1,"11,643.8",22.3 +Zee Entertainment Enterprises Ltd.,ZEEL,505537,MEDIA,BROADCASTING & CABLE TV,"2,509.6","2,105",332.8,13.65%,77.2,23.4,184.2,54.4,123,1.3,-102.2,-1.1 +eClerx Services Ltd.,ECLERX,532927,SOFTWARE & SERVICES,BPO/KPO,735.9,517,204.7,28.37%,30.3,6.1,182.4,46.3,136,28.2,506,105 +Sterlite Technologies Ltd.,STLTECH,532374,TELECOMMUNICATIONS EQUIPMENT,TELECOM CABLES,"1,497","1,281",213,14.26%,85,95,36,12,34,0.9,203,5.1 +HEG Ltd.,HEG,509631,GENERAL INDUSTRIALS,OTHER INDUSTRIAL GOODS,642.2,512.3,101.9,16.58%,38.5,8.5,82.9,21.7,96,24.9,439.5,113.9 +SBI Life Insurance Company Ltd.,SBILIFE,540719,BANKING AND FINANCE,LIFE INSURANCE,"28,816.2","28,183.8",609.9,2.12%,0,0,621.5,43.9,380.2,3.8,"1,842.2",18.4 +General Insurance Corporation of India,GICRE,540755,BANKING AND FINANCE,GENERAL INSURANCE,"13,465.9","11,574","1,464.6",11.20%,0,0,"1,855.4",243.7,"1,689",15.2,"6,628",37.8 +Tube Investments of India Ltd.,TIINDIA,540762,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"2,005.4","1,718.2",251.4,12.76%,34.6,7.7,244.8,63.4,181.4,9.4,717.5,37.1 +Honeywell Automation India Ltd.,HONAUT,517174,CONSUMER DURABLES,OTHER ELECTRICAL EQUIPMENT/PRODUCTS,"1,144.3",965.9,138.3,12.52%,13.8,0.7,163.9,42,121.9,137.8,443.4,503.9 +Indian Energy Exchange Ltd.,IEX,540750,BANKING AND FINANCE,EXCHANGE,133,16.6,92,84.73%,5.1,0.7,110.6,27.9,86.5,1,327.8,3.7 +ICICI Lombard General Insurance Company Ltd.,ICICIGI,540716,BANKING AND FINANCE,GENERAL INSURANCE,"5,271.1","4,612.4",743.5,14.16%,0,0,763.6,186.4,577.3,11.8,"1,757.1",35.8 +Aster DM Healthcare Ltd.,ASTERDM,540975,DIVERSIFIED CONSUMER SERVICES,HEALTHCARE FACILITIES,"3,325.2","2,939.4",377.3,11.38%,227.2,101.9,2.1,10.2,-30.8,-0.6,284.3,5.7 +Central Depository Services (India) Ltd.,CDSL,CDSL,OTHERS,INVESTMENT COMPANIES,230.1,77.9,129.4,62.40%,6.5,0,145.6,35.8,108.9,10.4,320.2,30.6 +Graphite India Ltd.,GRAPHITE,509488,GENERAL INDUSTRIALS,OTHER INDUSTRIAL GOODS,884,823,-30,-3.78%,19,4,992,190,804,41.1,856,43.9 +Grasim Industries Ltd.,GRASIM,500300,CEMENT AND CONSTRUCTION,CEMENT & CEMENT PRODUCTS,"30,505.3","25,995.9","4,224.8",13.98%,"1,245.2",397.8,"2,866.4",837.7,"1,163.8",17.7,"6,624.9",100.6 +KNR Constructions Ltd.,KNRCON,532942,CEMENT AND CONSTRUCTION,CONSTRUCTION & ENGINEERING,"1,043.8",806.9,231.6,22.30%,39.2,20.6,177.1,34.6,147.4,5.2,537.5,19.1 +Aditya Birla Capital Ltd.,ABCAPITAL,540691,DIVERSIFIED,HOLDING COMPANIES,"7,730.4","4,550.1","2,821.9",36.55%,48,"1,827",956.8,284.1,705,2.7,"5,231.9",20.1 +Dixon Technologies (India) Ltd.,DIXON,540699,CONSUMER DURABLES,CONSUMER ELECTRONICS,"4,943.9","4,744.3",198.9,4.02%,36.4,17.1,146.1,35.2,107.3,19,308.7,51.8 +Cholamandalam Financial Holdings Ltd.,CHOLAHLDNG,504973,DIVERSIFIED,HOLDING COMPANIES,"6,372.2","2,495.1","3,404.8",54.05%,52.1,"2,209.4","1,215.8",324.6,420.9,22.4,"1,532.3",81.6 +Cochin Shipyard Ltd.,COCHINSHIP,540678,TRANSPORTATION,MARINE PORT & SERVICES,"1,100.4",820.5,191.2,18.90%,18.9,9.6,251.4,69.9,181.5,13.8,429.9,32.7 +Bharat Dynamics Ltd.,BDL,541143,GENERAL INDUSTRIALS,DEFENCE,694.1,481.8,134,21.77%,17.4,0.8,194.1,47,147.1,8,425.4,23.2 +Lux Industries Ltd.,LUXIND,539542,TEXTILES APPARELS & ACCESSORIES,OTHER APPARELS & ACCESSORIES,643.6,584.2,55,8.61%,5.9,5.4,48,12.1,37.1,12.3,103.1,32.9 +Zensar Technologies Ltd.,ZENSARTECH,504067,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"1,277.1","1,009.9",230.9,18.61%,36.6,5.7,224.9,51,173.9,7.7,525.8,23.2 +PCBL Ltd.,PCBL,506590,CHEMICALS & PETROCHEMICALS,CARBON BLACK,"1,489.4","1,248.6",238.1,16.02%,48.2,21,171.6,48.8,122.6,3.2,431.6,11.4 +Zydus Wellness Ltd.,ZYDUSWELL,531335,FMCG,PACKAGED FOODS,444,423.1,16.8,3.82%,5.8,6.5,8.6,2.7,5.9,0.9,281.2,44.2 +Linde India Ltd.,LINDEINDIA,523457,GENERAL INDUSTRIALS,INDUSTRIAL GASES,729.9,537.7,173.6,24.41%,49.7,1.2,141.3,34.6,108.7,12.8,417.9,49 +FDC Ltd.,FDC,531599,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,513.6,409.9,76.4,15.71%,9.9,1.1,92.7,22.9,69.8,4.2,251.2,15.4 +The New India Assurance Company Ltd.,NIACL,540769,BANKING AND FINANCE,GENERAL INSURANCE,"10,571","10,773.4",-246.5,-2.33%,0,0,-242,-46.7,-176.1,-1.1,947,5.7 +Sundaram Finance Ltd.,SUNDARMFIN,590071,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),"1,710.6",322.5,"1,332.1",77.98%,43.6,820.3,470.6,142.8,365.4,33.2,"1,506.7",135.6 +TeamLease Services Ltd.,TEAMLEASE,539658,COMMERCIAL SERVICES & SUPPLIES,MISC. COMMERCIAL SERVICES,"2,285.6","2,240.8",31.8,1.40%,12.9,2.5,29.4,1.8,27.3,16.3,106.6,63.5 +Galaxy Surfactants Ltd.,GALAXYSURF,540935,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,985.8,858.2,124.9,12.70%,24.7,5.4,97.5,20.1,77.4,21.8,349.3,98.5 +Bandhan Bank Ltd.,BANDHANBNK,541153,BANKING AND FINANCE,BANKS,"5,032.2","1,400.2","1,583.4",35.25%,0,"2,048.6",947.2,226.1,721.2,4.5,"2,541.1",15.8 +ICICI Securities Ltd.,ISEC,541179,BANKING AND FINANCE,CAPITAL MARKETS,"1,249",433.5,810.2,64.87%,25.8,215.1,569.4,145.7,423.6,13.1,"1,238.1",38.3 +V-Mart Retail Ltd.,VMART,534976,RETAILING,DEPARTMENT STORES,551.4,548.8,0.7,0.12%,53.2,35.9,-86.4,-22.3,-64.1,-32.4,-103.1,-52.1 +Nippon Life India Asset Management Ltd.,NAM-INDIA,540767,BANKING AND FINANCE,ASSET MANAGEMENT COS.,475.4,156.1,241.4,60.73%,7.2,1.7,310.4,66.1,244.4,3.9,883.3,14.1 +Grindwell Norton Ltd.,GRINDWELL,506076,GENERAL INDUSTRIALS,OTHER INDUSTRIAL PRODUCTS,690,536,131.4,19.69%,16.9,1.8,135.3,33.1,101.9,9.2,378.3,34.2 +HDFC Life Insurance Company Ltd.,HDFCLIFE,540777,BANKING AND FINANCE,LIFE INSURANCE,"23,276.6","23,659.3",-508.1,-2.20%,0,0,-373.1,-657.5,378.2,1.8,"1,472.8",6.9 +Elgi Equipments Ltd.,ELGIEQUIP,522074,GENERAL INDUSTRIALS,INDUSTRIAL MACHINERY,817.8,663.4,142.7,17.71%,18.7,6.6,129.2,38.8,91.3,2.9,401.9,12.7 +Hindustan Aeronautics Ltd.,HAL,541154,GENERAL INDUSTRIALS,DEFENCE,"6,105.1","4,108.1","1,527.6",27.11%,349.6,0.3,"1,647",414.8,"1,236.7",18.5,"6,037.3",90.3 +BSE Ltd.,BSE,BSE,BANKING AND FINANCE,EXCHANGE,367,172.8,189.2,52.26%,22.7,8.5,163,63.6,120.5,8.8,706,52.1 +Rites Ltd.,RITES,541556,CEMENT AND CONSTRUCTION,CONSTRUCTION & ENGINEERING,608.8,444.5,137.8,23.67%,14.1,1.4,148.8,40.1,101.2,4.2,488.1,20.3 +Fortis Healthcare Ltd.,FORTIS,532843,DIVERSIFIED CONSUMER SERVICES,HEALTHCARE FACILITIES,"1,783.5","1,439.8",330.2,18.65%,84.1,31.8,231.4,48.8,173.7,2.3,547.6,7.3 +Varroc Engineering Ltd.,VARROC,541578,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"1,893.5","1,692.6",194.3,10.30%,84.9,50.3,65.9,18.2,54.2,3.5,146.5,9.6 +Adani Green Energy Ltd.,ADANIGREEN,ASM,UTILITIES,ELECTRIC UTILITIES,"2,589",521,"1,699",76.53%,474,"1,165",413,119,372,2.2,"1,305",8.2 +VIP Industries Ltd.,VIPIND,507880,TEXTILES APPARELS & ACCESSORIES,OTHER APPARELS & ACCESSORIES,548.7,493.2,52.9,9.68%,23.8,12.4,19.3,6,13.3,0.9,110.9,7.8 +CreditAccess Grameen Ltd.,CREDITACC,541770,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),"1,247.6",248.8,902.3,72.36%,12.3,423.9,466.8,119.7,347,21.8,"1,204.2",75.7 +CESC Ltd.,CESC,500084,UTILITIES,ELECTRIC UTILITIES,"4,414","3,706",646,14.84%,303,305,461,98,348,2.6,"1,447",10.9 +Jamna Auto Industries Ltd.,JAMNAAUTO,520051,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,608.7,528.2,79.1,13.03%,10.9,0.8,68.7,18.6,50.1,2.4,189.3,4.7 +Suprajit Engineering Ltd.,SUPRAJIT,532509,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,727.6,639.1,69.8,9.85%,25.7,13.6,49.2,14.5,34.8,2.5,146.9,10.6 +JK Paper Ltd.,JKPAPER,532162,COMMERCIAL SERVICES & SUPPLIES,PAPER & PAPER PRODUCTS,"1,708.8","1,242.8",407.3,24.68%,83.5,42,340.6,34.9,302.4,17.9,"1,220.6",72.1 +Bank of Maharashtra,MAHABANK,532525,BANKING AND FINANCE,BANKS,"5,735.5","1,179.4","1,920.5",37.90%,0,"2,635.7",935.7,16,919.8,1.3,"3,420.8",4.8 +Aavas Financiers Ltd.,AAVAS,541988,BANKING AND FINANCE,HOUSING FINANCE,497.6,123.5,367.8,74.03%,7.6,203.6,157.4,35.7,121.7,15.4,465.4,58.8 +HDFC Asset Management Company Ltd.,HDFCAMC,541729,BANKING AND FINANCE,ASSET MANAGEMENT COS.,765.4,162,481.1,74.81%,13,2.3,588.1,151.6,436.5,20.4,"1,659.3",77.7 +KEI Industries Ltd.,KEI,517569,CONSUMER DURABLES,OTHER ELECTRICAL EQUIPMENT/PRODUCTS,"1,954.2","1,742.7",203.9,10.47%,15.6,7.5,188.4,48.2,140.2,15.5,528.3,58.5 +Orient Electric Ltd.,ORIENTELEC,541301,CONSUMER DURABLES,CONSUMER ELECTRONICS,570.3,546.2,20.7,3.65%,14.2,5.2,23.4,4.9,18.4,0.9,95.3,4.5 +Deepak Nitrite Ltd.,DEEPAKNTR,506401,CHEMICALS & PETROCHEMICALS,COMMODITY CHEMICALS,"1,795.1","1,475.8",302.3,17.00%,39.4,2.7,277.2,72.1,205.1,15,797.9,58.5 +Fine Organic Industries Ltd.,FINEORG,541557,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,557.6,409.4,131.1,24.25%,14.4,0.7,133.1,28.9,103.4,33.7,458.8,149.6 +LTIMindtree Ltd.,LTIM,540005,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"9,048.6","7,274.1","1,631.3",18.32%,208.2,47,"1,519.3",357,"1,161.8",39.3,"4,427.5",149.6 +Dalmia Bharat Ltd.,DALBHARAT,542216,CEMENT AND CONSTRUCTION,CEMENT & CEMENT PRODUCTS,"3,234","2,56",589,18.70%,401,101,172,48,118,6.3,"1,041",54.8 +Godfrey Phillips India Ltd.,GODFRYPHLP,500163,FOOD BEVERAGES & TOBACCO,CIGARETTES-TOBACCO PRODUCTS,"1,412.5","1,151",223.6,16.27%,36.5,6.6,218.5,55.5,202.1,38.9,802.9,154.4 +Vaibhav Global Ltd.,VAIBHAVGBL,532156,TEXTILES APPARELS & ACCESSORIES,OTHER APPARELS & ACCESSORIES,708.4,641.5,63.5,9.01%,22.6,2.9,41.4,12.4,29.4,1.8,121.3,7.3 +Abbott India Ltd.,ABBOTINDIA,500488,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"1,549.7","1,113.3",380.9,25.49%,17.8,3.1,415.4,102.5,312.9,147.3,"1,081.4",508.9 +Adani Total Gas Ltd.,ATGL,ASM,UTILITIES,UTILITIES,"1,104.8",815.7,279.9,25.55%,37.6,27.3,224.2,57.2,172.7,1.6,571,5.2 +Nestle India Ltd.,NESTLEIND,500790,FMCG,PACKAGED FOODS,"5,070.1","3,811.9","1,224.9",24.32%,111.2,31.4,"1,222",313.9,908.1,94.2,"2,971.1",308.2 +Bayer Cropscience Ltd.,BAYERCROP,506285,CHEMICALS & PETROCHEMICALS,AGROCHEMICALS,"1,633.3","1,312.3",304.9,18.85%,11.6,3.7,305.7,82.8,222.9,49.6,844.4,188.1 +Amber Enterprises India Ltd.,AMBER,540902,CONSUMER DURABLES,CONSUMER ELECTRONICS,939.8,867.5,59.6,6.43%,45.2,36.6,-9.5,-3.8,-6.9,-2.1,156.8,46.5 +Rail Vikas Nigam Ltd.,RVNL,542649,CEMENT AND CONSTRUCTION,CONSTRUCTION & ENGINEERING,"5,210.3","4,616",298.3,6.07%,6.2,132.7,455.4,85.2,394.3,1.9,"1,478.8",7.1 +Metropolis Healthcare Ltd.,METROPOLIS,542650,DIVERSIFIED CONSUMER SERVICES,HEALTHCARE SERVICES,309.7,233.7,74.8,24.25%,22.2,5.7,48.1,12.5,35.5,6.9,133.4,26 +Polycab India Ltd.,POLYCAB,542652,CONSUMER DURABLES,OTHER ELECTRICAL EQUIPMENT/PRODUCTS,"4,253","3,608.8",608.9,14.44%,60.3,26.8,557.2,127.4,425.6,28.4,"1,607.2",107.1 +Multi Commodity Exchange of India Ltd.,MCX,534091,BANKING AND FINANCE,EXCHANGE,184,193.8,-28.7,-17.38%,6.6,0.1,-16.4,1.6,-19.1,-3.7,44.8,8.8 +IIFL Finance Ltd.,IIFL,532636,BANKING AND FINANCE,OTHER FINANCIAL SERVICES,"2,533.7",788.3,"1,600.8",64.66%,43.3,932.1,683.5,158,474.3,12.4,"1,690.7",44.4 +Ratnamani Metals & Tubes Ltd.,RATNAMANI,520111,METALS & MINING,IRON & STEEL/INTERM.PRODUCTS,"1,141.9",886.3,244.9,21.65%,23.6,10.8,221.1,56.8,163.9,23.4,622.6,88.8 +RHI Magnesita India Ltd.,RHIM,534076,GENERAL INDUSTRIALS,OTHER INDUSTRIAL GOODS,989.7,839,147.9,14.98%,44.2,8.5,97.9,26.3,71.3,3.5,-502.2,-24.3 +Birlasoft Ltd.,BSOFT,532400,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"1,325.4","1,102.7",207.1,15.81%,21.5,5.7,195.5,50.4,145.1,5.2,378.4,13.7 +EIH Ltd.,EIHOTEL,500840,HOTELS RESTAURANTS & TOURISM,HOTELS,552.5,387.6,142.9,26.94%,33.2,5.6,126.1,36.2,93.1,1.5,424.1,6.8 +Affle (India) Ltd.,AFFLE,542752,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,441.2,344.1,87.2,20.22%,18.4,5.5,73.2,6.4,66.8,5,264.3,19.8 +Westlife Foodworld Ltd.,WESTLIFE,505533,HOTELS RESTAURANTS & TOURISM,RESTAURANTS,618,516.5,98.2,15.98%,43.9,27.4,30.2,7.8,22.4,1.4,107.7,6.9 +IndiaMART InterMESH Ltd.,INDIAMART,542726,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,329.3,214.7,80,27.15%,8,2.3,104.3,23.9,69.4,11.4,321.1,53.6 +Infosys Ltd.,INFY,500209,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"39,626","29,554","9,44",24.21%,"1,166",138,"8,768","2,553","6,212",15,"24,871",60.1 +Sterling and Wilson Renewable Energy Ltd.,SWSOLAR,542760,COMMERCIAL SERVICES & SUPPLIES,CONSULTING SERVICES,776.7,758,1.5,0.19%,4.3,64.3,-50,4.6,-54.2,-2.9,-668.4,-35.2 +ABB India Ltd.,ABB,500002,GENERAL INDUSTRIALS,HEAVY ELECTRICAL EQUIPMENT,"2,846","2,330.7",438.5,15.84%,30.3,0.9,484.2,122.2,362.9,17.1,"1,208.7",57 +Poly Medicure Ltd.,POLYMED,531768,HEALTHCARE EQUIPMENT & SUPPLIES,HEALTHCARE SUPPLIES,351.4,253.1,84.2,24.97%,16,2.2,80.9,18.8,62.2,6.5,233.7,24.4 +GMM Pfaudler Ltd.,GMMPFAUDLR,505255,GENERAL INDUSTRIALS,INDUSTRIAL MACHINERY,946,795.5,142,15.15%,32.2,21.5,96.8,26.5,71.1,15.8,183.2,40.8 +Gujarat Fluorochemicals Ltd.,FLUOROCHEM,542812,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,960.3,783.7,163.1,17.23%,67.5,34.2,74.8,22.1,52.7,4.8,915.2,83.3 +360 One Wam Ltd.,360ONE,542772,BANKING AND FINANCE,OTHER FINANCIAL SERVICES,617.1,235.6,317.8,57.31%,13.7,139.9,226.8,40.8,186,5.2,696.8,19.5 +Tata Communications Ltd.,TATACOMM,500483,TELECOM SERVICES,OTHER TELECOM SERVICES,"4,897.9","3,857.1","1,015.5",20.84%,605.1,137.4,298.3,77.9,220.7,7.7,"1,322.3",46.4 +Alkyl Amines Chemicals Ltd.,ALKYLAMINE,506767,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,354.5,303.9,48.3,13.71%,12.5,1.7,36.4,9.2,27.2,5.3,171.3,33.5 +CSB Bank Ltd.,CSBBANK,542867,BANKING AND FINANCE,BANKS,835.8,317.5,174.6,25.41%,0,343.6,178,44.8,133.2,7.7,577.7,33.3 +Indian Railway Catering & Tourism Corporation Ltd.,IRCTC,542830,DIVERSIFIED CONSUMER SERVICES,TRAVEL SUPPORT SERVICES,"1,042.4",628.8,366.6,36.83%,14,4.4,395.2,100.5,294.7,3.7,"1,061.2",13.3 +Sumitomo Chemical India Ltd.,SUMICHEM,542920,CHEMICALS & PETROCHEMICALS,AGROCHEMICALS,928,715.5,187.9,20.80%,15.8,1.2,195.5,52,143.4,2.9,367.7,7.4 +Century Textiles & Industries Ltd.,CENTURYTEX,500040,COMMERCIAL SERVICES & SUPPLIES,PAPER & PAPER PRODUCTS,"1,114.9","1,069.2",33.8,3.07%,59.2,17,-30.5,-3.3,-30.4,-2.8,117.7,10.5 +SBI Cards and Payment Services Ltd.,SBICARD,543066,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),"4,221.4","2,018.8","1,327",32.47%,46.8,604.9,809.4,206.4,603,6.4,"2,302.2",24.3 +Hitachi Energy India Ltd.,POWERINDIA,543187,GENERAL INDUSTRIALS,HEAVY ELECTRICAL EQUIPMENT,"1,228.2","1,162.6",65.3,5.32%,22.5,10.7,32.4,7.6,24.7,5.8,82.5,19.5 +Suven Pharmaceuticals Ltd.,SUVENPHAR,543064,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,250.9,133.1,98,42.40%,11.9,0.5,105.4,25.8,79.6,3.1,431.8,17 +Tata Chemicals Ltd.,TATACHEM,500770,CHEMICALS & PETROCHEMICALS,COMMODITY CHEMICALS,"4,083","3,179",819,20.49%,234,145,627,120,428,16.8,"2,06",80.8 +Aarti Drugs Ltd.,AARTIDRUGS,524348,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,642.2,565.1,76.4,11.92%,12.6,8.2,56.3,16.7,39.6,4.3,180.2,19.6 +Gujarat Ambuja Exports Ltd.,GAEL,524226,FMCG,EDIBLE OILS,"1,157.7","1,012.2",103.3,9.26%,30.5,5.9,109.1,26.3,82.8,3.6,305.1,13.3 +Polyplex Corporation Ltd.,POLYPLEX,524051,COMMERCIAL SERVICES & SUPPLIES,CONTAINERS & PACKAGING,"1,595.7","1,451.5",120.6,7.67%,75.1,9.9,59.1,10.9,27.9,8.9,71.1,22.6 +Chalet Hotels Ltd.,CHALET,542399,HOTELS RESTAURANTS & TOURISM,HOTELS,318.2,188.6,126,40.04%,35,50.1,44.5,8,36.4,1.8,266.7,13 +Adani Enterprises Ltd.,ADANIENT,512599,COMMERCIAL SERVICES & SUPPLIES,COMMODITY TRADING & DISTRIBUTION,"23,066","20,087.2","2,430.1",10.79%,757,"1,342.8",791,397.8,227.8,2,"2,444.3",21.4 +YES Bank Ltd.,YESBANK,532648,BANKING AND FINANCE,BANKS,"7,980.6","2,377.1",810,12.06%,0,"4,793.6",304.4,75.7,228.6,0.1,836.6,0.3 +EPL Ltd.,EPL,500135,COMMERCIAL SERVICES & SUPPLIES,CONTAINERS & PACKAGING,"1,011.2",820.6,181,18.07%,83.6,30.6,76.4,25.4,50.5,1.6,251.9,7.9 +Network18 Media & Investments Ltd.,NETWORK18,532798,MEDIA,BROADCASTING & CABLE TV,"2,052.2","2,083.8",-218.3,-11.70%,56.8,66.2,-154.5,-6.5,-61,-0.6,-144.2,-1.4 +CIE Automotive India Ltd.,CIEINDIA,532756,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"2,299.4","1,934",345.4,15.15%,78.3,31,256.1,69.1,375.4,9.9,298.4,7.9 +Vedanta Ltd.,VEDL,500295,METALS & MINING,ALUMINIUM AND ALUMINIUM PRODUCTS,"39,585","27,466","11,479",29.47%,"2,642","2,523","8,177","9,092","-1,783",-4.8,"5,202",14 +Rossari Biotech Ltd.,ROSSARI,543213,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,484.8,419.9,63.6,13.15%,15.1,5,44.8,11.9,32.9,6,116.8,21.2 +KPIT Technologies Ltd.,KPITTECH,542651,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,"1,208.6",959.2,239.9,20.01%,48.1,13.6,187.7,46.3,140.9,5.2,486.9,18 +Intellect Design Arena Ltd.,INTELLECT,538835,SOFTWARE & SERVICES,IT SOFTWARE PRODUCTS,631.7,497.2,121.9,19.69%,33.7,0.8,96.5,25.7,70.4,5.2,316.6,23.2 +Balaji Amines Ltd.,BALAMINES,530999,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,387.3,326.8,53.8,14.13%,10.8,1.8,48,11.6,34.7,10.7,197.3,60.9 +UTI Asset Management Company Ltd.,UTIAMC,543238,BANKING AND FINANCE,ASSET MANAGEMENT COS.,405.6,172.5,231.5,57.30%,10.4,2.8,219.8,37,182.8,14.4,562.9,44.3 +Mazagon Dock Shipbuilders Ltd.,MAZDOCK,543237,TRANSPORTATION,SHIPPING,"2,079.2","1,651.1",176.6,9.66%,20.2,1.3,406.6,102.8,332.9,16.5,"1,327.6",65.8 +Computer Age Management Services Ltd.,CAMS,543232,BANKING AND FINANCE,CAPITAL MARKETS,284.7,153,122.1,44.39%,17.4,2,112.4,28.6,84.5,17.2,309.2,62.9 +Happiest Minds Technologies Ltd.,HAPPSTMNDS,543227,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,428.8,324,82.6,20.32%,14.6,11.2,79.1,20.7,58.5,3.9,232,15.6 +Triveni Turbine Ltd.,TRITURBINE,533655,GENERAL INDUSTRIALS,HEAVY ELECTRICAL EQUIPMENT,402.3,313.4,74.3,19.17%,5.1,0.6,83.2,19,64.2,2,233.1,7.3 +Angel One Ltd.,ANGELONE,ASM,BANKING AND FINANCE,CAPITAL MARKETS,"1,049.3",602.6,443.4,42.31%,11.2,26.4,407.2,102.7,304.5,36.3,"1,020.2",121.7 +Tanla Platforms Ltd.,TANLA,532790,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,"1,014.9",811.8,196.8,19.51%,22.6,1.8,178.7,36.2,142.5,10.6,514.7,38.3 +Max Healthcare Institute Ltd.,MAXHEALTH,543220,DIVERSIFIED CONSUMER SERVICES,HEALTHCARE FACILITIES,"1,408.6",975.8,387.4,28.42%,57.9,8.5,366.4,89.7,276.7,2.9,990.1,10.2 +Asahi India Glass Ltd.,ASAHIINDIA,515030,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"1,122.6",934,185.6,16.58%,43,34.4,111.3,30.2,86.9,3.6,343.5,14.1 +Prince Pipes & Fittings Ltd.,PRINCEPIPE,542907,GENERAL INDUSTRIALS,PLASTIC PRODUCTS,660.4,562.3,94.2,14.35%,22.5,0.7,92.8,22.2,70.6,5.2,219.8,19.9 +Route Mobile Ltd.,ROUTE,543228,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,"1,018.3",886.5,128.1,12.63%,21.4,6.5,103.8,15.5,88.8,14.2,365.3,58.3 +KPR Mill Ltd.,KPRMILL,532889,TEXTILES APPARELS & ACCESSORIES,TEXTILES,"1,533","1,212.9",298,19.72%,46,18.1,256,54.2,201.8,5.9,788.8,23.1 +Infibeam Avenues Ltd.,INFIBEAM,539807,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,792.6,719.7,70.2,8.89%,17.1,0.5,55.2,14.7,41,0.1,142.2,0.5 +Restaurant Brands Asia Ltd.,RBA,543248,HOTELS RESTAURANTS & TOURISM,RESTAURANTS,628.2,568.7,56.2,9.00%,78.6,31.5,-50.7,0,-46,-0.9,-220.3,-4.5 +Larsen & Toubro Ltd.,LT,500510,CEMENT AND CONSTRUCTION,CONSTRUCTION & ENGINEERING,"52,157","45,392.1","5,632",11.04%,909.9,864,"4,991.1","1,135.5","3,222.6",22.9,"12,255.3",89.2 +Gland Pharma Ltd.,GLAND,543245,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"1,426.6","1,049.3",324.1,23.60%,81.3,6,289.9,95.8,194.1,11.8,698.8,42.4 +Macrotech Developers Ltd.,LODHA,543287,REALTY,REALTY,"1,755.1","1,333.5",416.1,23.78%,29.3,123.1,269.2,62.4,201.9,2.1,"1,529.2",15.9 +Poonawalla Fincorp Ltd.,POONAWALLA,524000,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),745.3,178.9,531.7,71.98%,14.7,215.5,"1,124.6",270,860.2,11.2,"1,466.4",19.1 +The Fertilisers and Chemicals Travancore Ltd.,FACT,590024,FERTILIZERS,FERTILIZERS,"1,713.6","1,530.8",132.4,7.96%,5.3,61.2,105.2,0,105.2,1.6,508.4,7.9 +Home First Finance Company India Ltd.,HOMEFIRST,543259,BANKING AND FINANCE,HOUSING FINANCE,278,53.7,211.6,77.43%,2.8,117,96.4,22.1,74.3,8.4,266.2,30.2 +CG Power and Industrial Solutions Ltd.,CGPOWER,500093,GENERAL INDUSTRIALS,HEAVY ELECTRICAL EQUIPMENT,"2,019","1,692.9",308.6,15.42%,22.9,0.4,329.9,86.2,242.3,1.6,"1,1",7.2 +Laxmi Organic Industries Ltd.,LXCHEM,543277,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,660.5,613.3,38.9,5.97%,27.5,2.1,17.5,6.8,10.7,0.4,100.6,3.8 +Anupam Rasayan India Ltd.,ANURAS,543275,CHEMICALS & PETROCHEMICALS,AGROCHEMICALS,395.6,284.7,107.5,27.41%,19.8,20.4,70.7,22,40.7,3.8,178.9,16.6 +Kalyan Jewellers India Ltd.,KALYANKJIL,ASM,TEXTILES APPARELS & ACCESSORIES,GEMS & JEWELLERY,"4,427.7","4,100.9",313.7,7.11%,66.9,81.7,178.1,43.3,135.2,1.3,497.9,4.8 +Jubilant Pharmova Ltd.,JUBLPHARMA,530019,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"1,690.2","1,438.5",241.8,14.39%,96.6,66.1,89,35.9,62.5,3.9,-44.6,-2.8 +Indigo Paints Ltd.,INDIGOPNTS,543258,DIVERSIFIED CONSUMER SERVICES,FURNITURE-FURNISHING-PAINTS,273.4,228.7,41.8,15.45%,10,0.5,34.3,8.2,26.1,5.5,132.4,27.8 +Indian Railway Finance Corporation Ltd.,IRFC,543257,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),"6,767.5",33.5,"6,732.4",99.50%,2.1,"5,181.5","1,549.9",0,"1,549.9",1.2,"6,067.6",4.6 +Mastek Ltd.,MASTEK,523704,SOFTWARE & SERVICES,IT CONSULTING & SOFTWARE,770.4,642.5,123,16.07%,20.9,12.6,90.3,25,62.8,20.5,269.7,88 +Equitas Small Finance Bank Ltd.,EQUITASBNK,543243,BANKING AND FINANCE,BANKS,"1,540.4",616.8,330.2,24.30%,0,593.4,267,68.9,198.1,1.8,749.5,6.7 +Tata Teleservices (Maharashtra) Ltd.,TTML,532371,TELECOM SERVICES,TELECOM SERVICES,288.6,159.3,127.5,44.45%,36.3,403.2,-310.2,0,-310.2,-1.6,"-1,168.3",-6 +Praj Industries Ltd.,PRAJIND,522205,GENERAL INDUSTRIALS,INDUSTRIAL MACHINERY,893.3,798.4,84,9.52%,9.1,1,84.8,22.4,62.4,3.4,271.4,14.8 +Nazara Technologies Ltd.,NAZARA,543280,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,309.5,269.4,26.7,8.98%,15.1,2.7,21.2,-1.3,19.8,3,60,9.1 +Jubilant Ingrevia Ltd.,JUBLINGREA,543271,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,"1,028.5",902.3,117.7,11.54%,33.9,12.5,79.8,22.4,57.5,3.6,258.9,16.4 +Sona BLW Precision Forgings Ltd.,SONACOMS,543300,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,796.9,567.5,223.3,28.24%,53.4,6,164.1,40.1,123.8,2.1,462.8,7.9 +Chemplast Sanmar Ltd.,CHEMPLASTS,543336,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,"1,025",941.8,46,4.65%,35.3,38.6,9.2,-16.8,26.1,1.6,35.3,2.2 +Aptus Value Housing Finance India Ltd.,APTUS,543335,BANKING AND FINANCE,HOUSING FINANCE,344.5,50.6,277.5,83.18%,2.6,96.1,189.6,41.5,148,3,551.1,11.1 +Clean Science & Technology Ltd.,CLEAN,543318,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,187.1,106.3,74.8,41.32%,11.1,0.3,69.5,17.3,52.2,4.9,275.5,25.9 +Medplus Health Services Ltd.,MEDPLUS,543427,HEALTHCARE EQUIPMENT & SUPPLIES,HEALTHCARE SUPPLIES,"1,419","1,323.5",85.1,6.04%,55.5,23.5,16.4,1.9,14.6,1.2,58.3,4.9 +Nuvoco Vistas Corporation Ltd.,NUVOCO,543334,CEMENT AND CONSTRUCTION,CEMENT & CEMENT PRODUCTS,"2,578.9","2,243",329.9,12.82%,225.6,139.9,-29.6,-31.1,1.5,0,141.8,4 +Star Health and Allied Insurance Company Ltd.,STARHEALTH,543412,BANKING AND FINANCE,GENERAL INSURANCE,"3,463.2","3,295.8",165.7,4.79%,0,0,167.1,41.8,125.3,2.1,725.4,12.4 +Go Fashion (India) Ltd.,GOCOLORS,543401,TEXTILES APPARELS & ACCESSORIES,OTHER APPARELS & ACCESSORIES,192.8,132.2,56.6,29.98%,25.8,8.9,25.8,5.7,20,3.7,85.4,15.8 +PB Fintech Ltd.,POLICYBZR,543390,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,909.1,900.7,-89.1,-10.98%,22.3,7.2,-21.1,-0.3,-20.2,-0.5,-127.9,-2.8 +FSN E-Commerce Ventures Ltd.,NYKAA,543384,SOFTWARE & SERVICES,INTERNET & CATALOGUE RETAIL,"1,515.6","1,426.4",80.6,5.35%,54.6,21.3,13.3,4,5.8,0,19.8,0.1 +Krishna Institute of Medical Sciences Ltd.,KIMS,543308,DIVERSIFIED CONSUMER SERVICES,HEALTHCARE FACILITIES,655.4,475.2,177.3,27.17%,32.6,8.9,138.6,37.3,92,11.5,342.1,42.7 +Zomato Ltd.,ZOMATO,543320,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,"3,06","2,895",-47,-1.65%,128,16,21,-15,36,0,-496.8,-0.6 +Brightcom Group Ltd.,BCG,532368,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,"1,690.5","1,172.3",518,30.65%,72.3,0.1,445.8,124.3,321.5,1.6,"1,415.2",7 +Shyam Metalics and Energy Ltd.,SHYAMMETL,543299,METALS & MINING,IRON & STEEL/INTERM.PRODUCTS,"2,978.9","2,633.6",307.1,10.44%,176.5,35.4,133.4,-348.6,484.1,18.9,"1,049.9",41.2 +G R Infraprojects Ltd.,GRINFRA,543317,CEMENT AND CONSTRUCTION,ROADS & HIGHWAYS,"1,909.2","1,415.7",467.1,24.81%,61.7,144.6,287.1,69.9,217.2,22.5,"1,240.3",128.3 +RattanIndia Enterprises Ltd.,RTNINDIA,534597,UTILITIES,ELECTRIC UTILITIES,"1,618.1","1,392.8",1.5,0.11%,4.3,28.8,142.2,1.7,140.9,1,147.6,1.1 +Borosil Renewables Ltd.,BORORENEW,502219,CONSUMER DURABLES,HOUSEWARE,406.3,369.2,32.5,8.09%,31,9.6,28.9,-1.1,25.1,1.9,32.1,2.5 +HLE Glascoat Ltd.,HLEGLAS,522215,GENERAL INDUSTRIALS,INDUSTRIAL MACHINERY,227.8,198,26.5,11.79%,6.1,5.8,16.1,5.3,10,1.6,54.4,8 +Tata Investment Corporation Ltd.,TATAINVEST,501301,DIVERSIFIED,HOLDING COMPANIES,125,10.1,113.8,91.88%,0.2,4.7,110.1,-1.3,124.4,24.6,326.1,64.4 +Sapphire Foods India Ltd.,SAPPHIRE,543397,HOTELS RESTAURANTS & TOURISM,RESTAURANTS,650.1,527.5,115.1,17.91%,76.8,24.5,21.4,6.2,15.3,2.4,208.5,32.7 +Devyani International Ltd.,DEVYANI,543330,HOTELS RESTAURANTS & TOURISM,RESTAURANTS,826,665,154.4,18.84%,86.3,41.7,19,-16.8,33.4,0.3,177.5,1.5 +Vijaya Diagnostic Centre Ltd.,VIJAYA,543350,DIVERSIFIED CONSUMER SERVICES,HEALTHCARE SERVICES,145.6,81.5,57.4,41.31%,13.7,5.9,44.6,11,33.3,3.3,103.4,10.1 +C.E. Info Systems Ltd.,MAPMYINDIA,543425,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,99.3,50.1,41,44.98%,3.7,0.7,44.7,11.1,33,6.1,122.9,22.7 +Latent View Analytics Ltd.,LATENTVIEW,543398,SOFTWARE & SERVICES,DATA PROCESSING SERVICES,172.7,124.9,30.8,19.78%,2.3,0.8,44.7,10.6,34,1.7,153.6,7.5 +Metro Brands Ltd.,METROBRAND,543426,RETAILING,FOOTWEAR,571.9,400.3,155.4,27.96%,57.2,19.7,94.7,27.5,66.7,2.5,340,12.5 +Easy Trip Planners Ltd.,EASEMYTRIP,543272,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,144.6,76.9,64.8,45.71%,1,2,64.7,17.7,47.2,0.3,146,0.8 +Shree Renuka Sugars Ltd.,RENUKA,532670,FOOD BEVERAGES & TOBACCO,SUGAR,"2,564.7","2,491",63.7,2.49%,64.1,216.8,-207.2,-1.6,-204.9,-1,-286,-1.3 +One97 Communications Ltd.,PAYTM,543396,SOFTWARE & SERVICES,INTERNET SOFTWARE & SERVICES,"2,662.5","2,749.6",-231,-9.17%,180.1,7,-279.9,12.7,-290.5,-5,"-1,207.9",-19 +MTAR Technologies Ltd.,MTARTECH,543270,GENERAL INDUSTRIALS,DEFENCE,167.7,130.7,36.1,21.64%,5.8,5.5,25.7,5.2,20.5,6.7,103.3,33.6 +Capri Global Capital Ltd.,CGCL,531595,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),557.4,229.3,304.8,54.70%,23.1,195.8,86,20.8,65.2,3.2,231.2,11.2 +GMR Airports Infrastructure Ltd.,GMRINFRA,ASM,CEMENT AND CONSTRUCTION,CONSTRUCTION & ENGINEERING,"2,185","1,336.8",726.7,35.22%,373,695.8,-252,54.9,-91,-0.1,-370.9,-0.6 +Triveni Engineering & Industries Ltd.,TRIVENI,532356,FOOD BEVERAGES & TOBACCO,SUGAR,"1,629.7","1,554.5",62.9,3.89%,25.8,10.2,39.3,10.1,29.1,1.3,434.3,19.8 +Delhivery Ltd.,DELHIVERY,543529,TRANSPORTATION,TRANSPORTATION - LOGISTICS,"2,043","1,957.3",-15.6,-0.80%,171.2,19.6,-105.2,-2.1,-102.9,-1.4,-546.7,-7.5 +Life Insurance Corporation of India,LICI,543526,BANKING AND FINANCE,LIFE INSURANCE,"202,394.9","193,612.5","8,445",4.18%,0,0,"8,696.5","1,083.9","8,030.3",12.7,"37,204.8",58.8 +Campus Activewear Ltd.,CAMPUS,543523,RETAILING,FOOTWEAR,259.1,234.2,24.5,9.46%,18.1,6.5,0.4,0.1,0.3,0,103.1,3.4 +Motherson Sumi Wiring India Ltd.,MSUMI,543498,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"2,110.2","1,856.5",248.1,11.79%,36.4,7.4,210,54.1,155.9,0.3,523.6,1.2 +Olectra Greentech Ltd.,OLECTRA,532439,AUTOMOBILES & AUTO COMPONENTS,COMMERCIAL VEHICLES,310.3,266.6,40.5,13.20%,8.8,9.7,25.2,8,18.6,2.2,78.5,9.6 +Patanjali Foods Ltd.,PATANJALI,500368,FMCG,EDIBLE OILS,"7,845.8","7,426.6",395.3,5.05%,60.1,24,335.1,80.5,254.5,7,875.2,24.2 +Raymond Ltd.,RAYMOND,500330,TEXTILES APPARELS & ACCESSORIES,TEXTILES,"2,320.7","1,938.8",314.6,13.96%,65.4,89.3,204.2,50.7,159.8,24,"1,514.2",227.5 +Swan Energy Ltd.,SWANENERGY,503310,REALTY,REALTY,"1,230.1",966.3,257,21.01%,27.1,58.3,178.4,12.8,84.6,6.7,308.4,11.7 +Samvardhana Motherson International Ltd.,MOTHERSON,517334,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"23,639.2","21,585","1,888.8",8.05%,867.4,487.9,449.5,229.2,201.6,0.3,"1,910.3",2.8 +Vedant Fashions Ltd.,MANYAVAR,543463,RETAILING,SPECIALTY RETAIL,233.4,125.5,92.8,42.51%,32.5,10.7,64.8,16.1,48.7,2,399.9,16.5 +Adani Wilmar Ltd.,AWL,543458,FMCG,EDIBLE OILS,"12,331.2","12,123.5",143.7,1.17%,95.7,220.2,-161.8,-31.5,-130.7,-1,130.1,1 +Mahindra Lifespace Developers Ltd.,MAHLIFE,532313,REALTY,REALTY,25.7,52.7,-34.9,-196.45%,3.1,0.2,-30.3,-10.8,-18.9,-1.2,10.5,0.7 +Tejas Networks Ltd.,TEJASNET,540595,TELECOM SERVICES,OTHER TELECOM SERVICES,413.9,383,13,3.28%,41.7,7,-17.7,-5.1,-12.6,-0.7,-61.3,-3.5 +Aether Industries Ltd.,AETHER,543534,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,178.3,118.2,46,28.00%,9.7,1.6,48.7,12.1,36.7,2.8,139.1,10.5 +JBM Auto Ltd.,JBMA,ASM,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"1,238.8","1,091.3",139.7,11.35%,41.2,47.9,58.3,11.3,44.2,3.7,136.8,11.6 +Deepak Fertilisers & Petrochemicals Corporation Ltd.,DEEPAKFERT,500645,CHEMICALS & PETROCHEMICALS,COMMODITY CHEMICALS,"2,443.2","2,138.1",286.1,11.80%,81.2,107.1,116.8,53.3,60.1,4.8,674.5,53.4 +Sharda Cropchem Ltd.,SHARDACROP,538666,CHEMICALS & PETROCHEMICALS,AGROCHEMICALS,604.3,559.6,21.2,3.65%,74,4.6,-33.8,-6.3,-27.6,-3.1,191,21.2 +Shoppers Stop Ltd.,SHOPERSTOP,532638,RETAILING,DEPARTMENT STORES,"1,049.7",878.2,160.9,15.49%,108.2,54.9,3.5,0.8,2.7,0.2,94.2,8.6 +BEML Ltd.,BEML,500048,AUTOMOBILES & AUTO COMPONENTS,COMMERCIAL VEHICLES,924,855.3,61.5,6.70%,15.8,10.8,42.2,-9.6,51.8,12.4,200.8,48.2 +Lemon Tree Hotels Ltd.,LEMONTREE,541233,HOTELS RESTAURANTS & TOURISM,HOTELS,230.1,125.3,101.9,44.84%,22.6,47.3,34.8,8.6,22.6,0.3,130.1,1.6 +Rainbow Childrens Medicare Ltd.,RAINBOW,543524,DIVERSIFIED CONSUMER SERVICES,HEALTHCARE FACILITIES,340.5,215.1,117.6,35.34%,26.8,13.3,85.2,22.1,62.9,6.2,215.4,21.2 +UCO Bank,UCOBANK,532505,BANKING AND FINANCE,BANKS,"5,865.6","1,581.5",981.9,18.81%,0,"3,302.3",639.8,238.1,403.5,0.3,"1,84",1.5 +Piramal Pharma Ltd.,PPLPHARMA,543635,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"1,960.6","1,645.7",265.6,13.90%,184.5,109.9,20.4,34.5,5,0,-133.6,-1 +KSB Ltd.,KSB,500249,GENERAL INDUSTRIALS,INDUSTRIAL MACHINERY,572.2,493.4,70.3,12.47%,12.3,2,64.5,17.1,50.1,14.4,209.7,60.3 +Data Patterns (India) Ltd.,DATAPATTNS,543428,GENERAL INDUSTRIALS,DEFENCE,119.2,67.5,40.8,37.63%,3.1,2.3,46.3,12.5,33.8,6,148.3,26.5 +Global Health Ltd.,MEDANTA,543654,DIVERSIFIED CONSUMER SERVICES,HEALTHCARE FACILITIES,864.7,631.1,212.9,25.22%,42.9,20.1,170.6,45.4,125.2,4.7,408.9,15.2 +Aarti Industries Ltd.,AARTIIND,524208,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,"1,454","1,221.2",232.8,16.01%,93,58.2,81.6,-9.1,90.7,2.5,446.2,12.3 +BLS International Services Ltd.,BLS,540073,DIVERSIFIED CONSUMER SERVICES,TRAVEL SUPPORT SERVICES,416.4,321,86.7,21.27%,7.3,1,87.2,5.2,78.7,1.9,267.6,6.5 +Archean Chemical Industries Ltd.,ACI,543657,CHEMICALS & PETROCHEMICALS,COMMODITY CHEMICALS,301.7,195,95.5,32.86%,17.5,1.9,87.3,21.3,66,5.4,394.4,32.1 +Adani Power Ltd.,ADANIPOWER,ASM,UTILITIES,ELECTRIC UTILITIES,"14,935.7","7,819.2","5,171.4",39.81%,"1,004.5",888.4,"5,223.6","-1,370.6","6,594.2",16.5,"20,604.8",53.4 +Craftsman Automation Ltd.,CRAFTSMAN,543276,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"1,183.8",941.6,237.5,20.14%,66.8,41.6,133.8,29.6,94.5,44.1,298.3,141.2 +NMDC Ltd.,NMDC,526371,METALS & MINING,MINING,"4,335","2,823.6","1,190.4",29.66%,88.8,18.6,"1,404.1",379,"1,026.2",3.5,"5,862.2",20 +Epigral Ltd.,EPIGRAL,543332,CHEMICALS & PETROCHEMICALS,SPECIALTY CHEMICALS,479.1,370.2,107.9,22.57%,31.5,21.3,56.1,17.9,38,9.1,223.4,53.8 +Apar Industries Ltd.,APARINDS,532259,CONSUMER DURABLES,OTHER ELECTRICAL EQUIPMENT/PRODUCTS,"3,944.7","3,576.2",349.8,8.91%,28.2,103.1,237.3,62.9,173.9,45.4,783.9,204.8 +Bikaji Foods International Ltd.,BIKAJI,543653,FMCG,PACKAGED FOODS,614.7,521,87.7,14.41%,15.6,2.9,75.2,15.4,61.2,2.5,173.6,6.9 +Five-Star Business Finance Ltd.,FIVESTAR,543663,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),522.4,133.2,375,72.28%,5.7,105.9,267,67.6,199.4,6.8,703,24.1 +Ingersoll-Rand (India) Ltd.,INGERRAND,500210,GENERAL INDUSTRIALS,INDUSTRIAL MACHINERY,282.8,210.7,65.7,23.76%,4.6,0.6,67,17.2,49.7,15.8,218.5,69.2 +KFIN Technologies Ltd.,KFINTECH,543720,BANKING AND FINANCE,OTHER FINANCIAL SERVICES,215.3,115.3,93.7,44.82%,12.6,3.2,84.2,22.3,61.4,3.6,215.1,12.6 +Piramal Enterprises Ltd.,PEL,500302,BANKING AND FINANCE,FINANCE (INCLUDING NBFCS),"2,205.2","1,320.1","1,117.9",50.97%,38.3,"1,038.9",-11.8,10.7,48.2,2,"3,906.5",173.9 +NMDC Steel Ltd.,NSLNISP,543768,METALS & MINING,IRON & STEEL/INTERM.PRODUCTS,290.3,349.6,-72.2,-26.04%,74.5,40.8,-174.7,-43.6,-131.1,-0.5,, +Eris Lifesciences Ltd.,ERIS,540596,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,508.8,324.2,181.1,35.85%,42.1,16.3,126.2,3.9,123.4,9.1,385.6,28.3 +Mankind Pharma Ltd.,MANKIND,543904,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,"2,768.1","2,025.5",682.6,25.21%,96.5,8.6,637.5,129.8,501,12.5,"1,564.8",39.1 +Kaynes Technology India Ltd.,KAYNES,ASM,CONSUMER DURABLES,OTHER ELECTRICAL EQUIPMENT/PRODUCTS,369.8,312.1,48.8,13.52%,6.5,11.8,39.4,7.1,32.3,5.5,143.2,24.6 +Safari Industries (India) Ltd.,SAFARI,523025,TEXTILES APPARELS & ACCESSORIES,OTHER APPARELS & ACCESSORIES,372.9,306.6,63.5,17.15%,12.2,2.2,51.9,12.1,39.8,16.7,162.3,68.2 +Saregama India Ltd.,SAREGAMA,532163,MEDIA,MOVIES & ENTERTAINMENT,185.6,111.5,60.9,35.32%,8.2,0.2,65.6,17.6,48.1,2.5,193.4,10 +Syrma SGS Technology Ltd.,SYRMA,543573,CONSUMER DURABLES,OTHER ELECTRICAL EQUIPMENT/PRODUCTS,720.6,662.7,49,6.88%,11.6,8,37,6.4,28.3,1.6,132.4,7.5 +Jindal Saw Ltd.,JINDALSAW,ASM,GENERAL INDUSTRIALS,OTHER INDUSTRIAL PRODUCTS,"5,488.9","4,662",804.2,14.71%,142.5,188.7,495.6,139.6,375.7,11.8,"1,135.8",35.5 +Godawari Power & Ispat Ltd.,GPIL,532734,METALS & MINING,IRON & STEEL/INTERM.PRODUCTS,"1,314.2",929.6,361.4,28.00%,34.8,10.2,339.6,86.1,256.9,20.6,785.5,63 +Gillette India Ltd.,GILLETTE,507815,FMCG,PERSONAL PRODUCTS,676.2,530.8,136.7,20.48%,20.1,0.1,125.2,32.5,92.7,28.4,361.6,111 +Symphony Ltd.,SYMPHONY,517385,CONSUMER DURABLES,CONSUMER ELECTRONICS,286,234,41,14.91%,7,2,43,8,35,5.1,114,16.5 +Glenmark Life Sciences Ltd.,GLS,543322,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,600.7,428.3,167.1,28.06%,13.1,0.4,158.9,40.2,118.7,9.7,505.5,41.3 +Usha Martin Ltd.,USHAMART,517146,METALS & MINING,IRON & STEEL PRODUCTS,806,640.4,144.3,18.39%,18,6.4,141.2,35,109.5,3.6,399.4,13.1 +Ircon International Ltd.,IRCON,541956,CEMENT AND CONSTRUCTION,CONSTRUCTION & ENGINEERING,"3,136.3","2,771.2",215.7,7.22%,27.1,36.9,301.2,77.6,250.7,2.7,884.6,9.4 +Ujjivan Small Finance Bank Ltd.,UJJIVANSFB,542904,BANKING AND FINANCE,BANKS,"1,579.8",528.6,483.4,34.75%,0,567.8,436.4,108.7,327.7,1.7,"1,254.5",6.4 +Procter & Gamble Health Ltd.,PGHL,500126,PHARMACEUTICALS & BIOTECHNOLOGY,PHARMACEUTICALS,311,216.3,88.7,29.08%,6.5,0.2,88,22.5,65.6,39.5,231.4,139.4 +Allcargo Logistics Ltd.,ALLCARGO,532749,TRANSPORTATION,TRANSPORTATION - LOGISTICS,"3,336.3","3,188.8",118,3.57%,106.7,36.7,14.2,1.3,21.8,0.9,361.9,14.7 +Sheela Foam Ltd.,SFL,540203,DIVERSIFIED CONSUMER SERVICES,FURNITURE-FURNISHING-PAINTS,637.6,547,66.2,10.80%,21.9,8.6,60.2,15.6,44,4.5,192.4,17.7 +Alok Industries Ltd.,ALOKINDS,521070,TEXTILES APPARELS & ACCESSORIES,TEXTILES,"1,369.3","1,323.1",35.9,2.64%,78.6,142.2,-174.6,0,-174.8,-0.3,-948.4,-1.9 +Minda Corporation Ltd.,MINDACORP,538962,AUTOMOBILES & AUTO COMPONENTS,AUTO PARTS & EQUIPMENT,"1,197.9","1,064.5",131.3,10.98%,41.4,14.9,77,18.7,58.8,2.5,278.2,11.6 +Concord Biotech Ltd.,CONCORDBIO,543960,PHARMACEUTICALS & BIOTECHNOLOGY,BIOTECHNOLOGY,270.5,143.2,119.2,45.43%,13.3,0.8,113.2,28.7,81,7.7,, \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/extracting-results-with-an-agent.ipynb b/python/docs/src/user-guide/core-user-guide/cookbook/extracting-results-with-an-agent.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/extracting-results-with-an-agent.ipynb rename to python/docs/src/user-guide/core-user-guide/cookbook/extracting-results-with-an-agent.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/index.md b/python/docs/src/user-guide/core-user-guide/cookbook/index.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/index.md rename to python/docs/src/user-guide/core-user-guide/cookbook/index.md diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/instrumenting.md b/python/docs/src/user-guide/core-user-guide/cookbook/instrumenting.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/instrumenting.md rename to python/docs/src/user-guide/core-user-guide/cookbook/instrumenting.md diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/langgraph-agent.ipynb b/python/docs/src/user-guide/core-user-guide/cookbook/langgraph-agent.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/langgraph-agent.ipynb rename to python/docs/src/user-guide/core-user-guide/cookbook/langgraph-agent.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/llamaindex-agent.ipynb b/python/docs/src/user-guide/core-user-guide/cookbook/llamaindex-agent.ipynb similarity index 99% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/llamaindex-agent.ipynb rename to python/docs/src/user-guide/core-user-guide/cookbook/llamaindex-agent.ipynb index e75293fcec8a..03894e68699e 100644 --- a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/llamaindex-agent.ipynb +++ b/python/docs/src/user-guide/core-user-guide/cookbook/llamaindex-agent.ipynb @@ -38,7 +38,6 @@ "outputs": [], "source": [ "import os\n", - "from dataclasses import dataclass\n", "from typing import List, Optional\n", "\n", "from autogen_core import AgentId, MessageContext, RoutedAgent, SingleThreadedAgentRuntime, message_handler\n", @@ -57,7 +56,8 @@ "from llama_index.embeddings.openai import OpenAIEmbedding\n", "from llama_index.llms.azure_openai import AzureOpenAI\n", "from llama_index.llms.openai import OpenAI\n", - "from llama_index.tools.wikipedia import WikipediaToolSpec" + "from llama_index.tools.wikipedia import WikipediaToolSpec\n", + "from pydantic import BaseModel" ] }, { @@ -73,15 +73,13 @@ "metadata": {}, "outputs": [], "source": [ - "@dataclass\n", - "class Resource:\n", + "class Resource(BaseModel):\n", " content: str\n", " node_id: str\n", " score: Optional[float] = None\n", "\n", "\n", - "@dataclass\n", - "class Message:\n", + "class Message(BaseModel):\n", " content: str\n", " sources: Optional[List[Resource]] = None" ] diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/llm-usage-logger.ipynb b/python/docs/src/user-guide/core-user-guide/cookbook/llm-usage-logger.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/llm-usage-logger.ipynb rename to python/docs/src/user-guide/core-user-guide/cookbook/llm-usage-logger.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/local-llms-ollama-litellm.ipynb b/python/docs/src/user-guide/core-user-guide/cookbook/local-llms-ollama-litellm.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/local-llms-ollama-litellm.ipynb rename to python/docs/src/user-guide/core-user-guide/cookbook/local-llms-ollama-litellm.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/openai-assistant-agent.ipynb b/python/docs/src/user-guide/core-user-guide/cookbook/openai-assistant-agent.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/openai-assistant-agent.ipynb rename to python/docs/src/user-guide/core-user-guide/cookbook/openai-assistant-agent.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/structured-output-agent.ipynb b/python/docs/src/user-guide/core-user-guide/cookbook/structured-output-agent.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/structured-output-agent.ipynb rename to python/docs/src/user-guide/core-user-guide/cookbook/structured-output-agent.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/termination-with-intervention.ipynb b/python/docs/src/user-guide/core-user-guide/cookbook/termination-with-intervention.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/termination-with-intervention.ipynb rename to python/docs/src/user-guide/core-user-guide/cookbook/termination-with-intervention.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/tool-use-with-intervention.ipynb b/python/docs/src/user-guide/core-user-guide/cookbook/tool-use-with-intervention.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/tool-use-with-intervention.ipynb rename to python/docs/src/user-guide/core-user-guide/cookbook/tool-use-with-intervention.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/topic-subscription-scenarios.ipynb b/python/docs/src/user-guide/core-user-guide/cookbook/topic-subscription-scenarios.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/topic-subscription-scenarios.ipynb rename to python/docs/src/user-guide/core-user-guide/cookbook/topic-subscription-scenarios.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/agent-and-multi-agent-application.md b/python/docs/src/user-guide/core-user-guide/core-concepts/agent-and-multi-agent-application.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/agent-and-multi-agent-application.md rename to python/docs/src/user-guide/core-user-guide/core-concepts/agent-and-multi-agent-application.md diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/agent-identity-and-lifecycle.md b/python/docs/src/user-guide/core-user-guide/core-concepts/agent-identity-and-lifecycle.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/agent-identity-and-lifecycle.md rename to python/docs/src/user-guide/core-user-guide/core-concepts/agent-identity-and-lifecycle.md diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/agent-lifecycle.svg b/python/docs/src/user-guide/core-user-guide/core-concepts/agent-lifecycle.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/agent-lifecycle.svg rename to python/docs/src/user-guide/core-user-guide/core-concepts/agent-lifecycle.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/application-stack.md b/python/docs/src/user-guide/core-user-guide/core-concepts/application-stack.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/application-stack.md rename to python/docs/src/user-guide/core-user-guide/core-concepts/application-stack.md diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/application-stack.svg b/python/docs/src/user-guide/core-user-guide/core-concepts/application-stack.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/application-stack.svg rename to python/docs/src/user-guide/core-user-guide/core-concepts/application-stack.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/architecture-distributed.svg b/python/docs/src/user-guide/core-user-guide/core-concepts/architecture-distributed.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/architecture-distributed.svg rename to python/docs/src/user-guide/core-user-guide/core-concepts/architecture-distributed.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/architecture-standalone.svg b/python/docs/src/user-guide/core-user-guide/core-concepts/architecture-standalone.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/architecture-standalone.svg rename to python/docs/src/user-guide/core-user-guide/core-concepts/architecture-standalone.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/architecture.md b/python/docs/src/user-guide/core-user-guide/core-concepts/architecture.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/architecture.md rename to python/docs/src/user-guide/core-user-guide/core-concepts/architecture.md diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/code-gen-example.svg b/python/docs/src/user-guide/core-user-guide/core-concepts/code-gen-example.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/code-gen-example.svg rename to python/docs/src/user-guide/core-user-guide/core-concepts/code-gen-example.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/subscription.svg b/python/docs/src/user-guide/core-user-guide/core-concepts/subscription.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/subscription.svg rename to python/docs/src/user-guide/core-user-guide/core-concepts/subscription.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/topic-and-subscription.md b/python/docs/src/user-guide/core-user-guide/core-concepts/topic-and-subscription.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/topic-and-subscription.md rename to python/docs/src/user-guide/core-user-guide/core-concepts/topic-and-subscription.md diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/type-subscription-multi-tenant.svg b/python/docs/src/user-guide/core-user-guide/core-concepts/type-subscription-multi-tenant.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/type-subscription-multi-tenant.svg rename to python/docs/src/user-guide/core-user-guide/core-concepts/type-subscription-multi-tenant.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/type-subscription-single-tenant-multiple-topics.svg b/python/docs/src/user-guide/core-user-guide/core-concepts/type-subscription-single-tenant-multiple-topics.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/type-subscription-single-tenant-multiple-topics.svg rename to python/docs/src/user-guide/core-user-guide/core-concepts/type-subscription-single-tenant-multiple-topics.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/type-subscription-single-tenant-single-topic.svg b/python/docs/src/user-guide/core-user-guide/core-concepts/type-subscription-single-tenant-single-topic.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/core-concepts/type-subscription-single-tenant-single-topic.svg rename to python/docs/src/user-guide/core-user-guide/core-concepts/type-subscription-single-tenant-single-topic.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/code-execution-groupchat.ipynb b/python/docs/src/user-guide/core-user-guide/design-patterns/code-execution-groupchat.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/code-execution-groupchat.ipynb rename to python/docs/src/user-guide/core-user-guide/design-patterns/code-execution-groupchat.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/coder-reviewer-data-flow.svg b/python/docs/src/user-guide/core-user-guide/design-patterns/coder-reviewer-data-flow.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/coder-reviewer-data-flow.svg rename to python/docs/src/user-guide/core-user-guide/design-patterns/coder-reviewer-data-flow.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/concurrent-agents.ipynb b/python/docs/src/user-guide/core-user-guide/design-patterns/concurrent-agents.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/concurrent-agents.ipynb rename to python/docs/src/user-guide/core-user-guide/design-patterns/concurrent-agents.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/group-chat.ipynb b/python/docs/src/user-guide/core-user-guide/design-patterns/group-chat.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/group-chat.ipynb rename to python/docs/src/user-guide/core-user-guide/design-patterns/group-chat.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/groupchat.svg b/python/docs/src/user-guide/core-user-guide/design-patterns/groupchat.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/groupchat.svg rename to python/docs/src/user-guide/core-user-guide/design-patterns/groupchat.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/handoffs.ipynb b/python/docs/src/user-guide/core-user-guide/design-patterns/handoffs.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/handoffs.ipynb rename to python/docs/src/user-guide/core-user-guide/design-patterns/handoffs.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/handoffs.svg b/python/docs/src/user-guide/core-user-guide/design-patterns/handoffs.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/handoffs.svg rename to python/docs/src/user-guide/core-user-guide/design-patterns/handoffs.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/intro.md b/python/docs/src/user-guide/core-user-guide/design-patterns/intro.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/intro.md rename to python/docs/src/user-guide/core-user-guide/design-patterns/intro.md diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/mixture-of-agents.ipynb b/python/docs/src/user-guide/core-user-guide/design-patterns/mixture-of-agents.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/mixture-of-agents.ipynb rename to python/docs/src/user-guide/core-user-guide/design-patterns/mixture-of-agents.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/multi-agent-debate.ipynb b/python/docs/src/user-guide/core-user-guide/design-patterns/multi-agent-debate.ipynb similarity index 99% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/multi-agent-debate.ipynb rename to python/docs/src/user-guide/core-user-guide/design-patterns/multi-agent-debate.ipynb index 483256977092..08f907eb25a9 100644 --- a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/multi-agent-debate.ipynb +++ b/python/docs/src/user-guide/core-user-guide/design-patterns/multi-agent-debate.ipynb @@ -291,7 +291,7 @@ "A --- B\n", "| |\n", "| |\n", - "C --- D\n", + "D --- C\n", "```\n", "\n", "Each solver agent is connected to two other solver agents. \n", diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/reflection.ipynb b/python/docs/src/user-guide/core-user-guide/design-patterns/reflection.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/reflection.ipynb rename to python/docs/src/user-guide/core-user-guide/design-patterns/reflection.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/sequential-workflow.ipynb b/python/docs/src/user-guide/core-user-guide/design-patterns/sequential-workflow.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/sequential-workflow.ipynb rename to python/docs/src/user-guide/core-user-guide/design-patterns/sequential-workflow.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/sequential-workflow.svg b/python/docs/src/user-guide/core-user-guide/design-patterns/sequential-workflow.svg similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/sequential-workflow.svg rename to python/docs/src/user-guide/core-user-guide/design-patterns/sequential-workflow.svg diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/faqs.md b/python/docs/src/user-guide/core-user-guide/faqs.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/faqs.md rename to python/docs/src/user-guide/core-user-guide/faqs.md diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/agent-and-agent-runtime.ipynb b/python/docs/src/user-guide/core-user-guide/framework/agent-and-agent-runtime.ipynb similarity index 99% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/agent-and-agent-runtime.ipynb rename to python/docs/src/user-guide/core-user-guide/framework/agent-and-agent-runtime.ipynb index e8e34a73b420..7f18ac2bce3e 100644 --- a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/agent-and-agent-runtime.ipynb +++ b/python/docs/src/user-guide/core-user-guide/framework/agent-and-agent-runtime.ipynb @@ -133,7 +133,7 @@ " response = await self._delegate.on_messages(\n", " [TextMessage(content=message.content, source=\"user\")], ctx.cancellation_token\n", " )\n", - " print(f\"{self.id.type} responded: {response.chat_message.content}\")" + " print(f\"{self.id.type} responded: {response.chat_message}\")" ] }, { diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/component-config.ipynb b/python/docs/src/user-guide/core-user-guide/framework/component-config.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/component-config.ipynb rename to python/docs/src/user-guide/core-user-guide/framework/component-config.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/distributed-agent-runtime.ipynb b/python/docs/src/user-guide/core-user-guide/framework/distributed-agent-runtime.ipynb similarity index 97% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/distributed-agent-runtime.ipynb rename to python/docs/src/user-guide/core-user-guide/framework/distributed-agent-runtime.ipynb index ce11bc2d4960..22b1e8b23d9d 100644 --- a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/distributed-agent-runtime.ipynb +++ b/python/docs/src/user-guide/core-user-guide/framework/distributed-agent-runtime.ipynb @@ -99,7 +99,7 @@ "All agents publish and subscribe to the default topic, so they can see all\n", "messages being published.\n", "\n", - "To run the agents, we publishes a message from a worker." + "To run the agents, we publish a message from a worker." ] }, { @@ -189,16 +189,21 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Cross-Language Runtimes\n", + "## Cross-Language Runtimes\n", "The process described above is largely the same, however all message types MUST use shared protobuf schemas for all cross-agent message types.\n", "\n", - "# Next Steps\n", + "## Next Steps\n", "To see complete examples of using distributed runtime, please take a look at the following samples:\n", "\n", "- [Distributed Workers](https://github.com/microsoft/autogen/tree/main/python/samples/core_grpc_worker_runtime) \n", "- [Distributed Semantic Router](https://github.com/microsoft/autogen/tree/main/python/samples/core_semantic_router) \n", "- [Distributed Group Chat](https://github.com/microsoft/autogen/tree/main/python/samples/core_distributed-group-chat) \n" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [] } ], "metadata": { diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/logging.md b/python/docs/src/user-guide/core-user-guide/framework/logging.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/logging.md rename to python/docs/src/user-guide/core-user-guide/framework/logging.md diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/message-and-communication.ipynb b/python/docs/src/user-guide/core-user-guide/framework/message-and-communication.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/message-and-communication.ipynb rename to python/docs/src/user-guide/core-user-guide/framework/message-and-communication.ipynb diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/telemetry.md b/python/docs/src/user-guide/core-user-guide/framework/telemetry.md similarity index 68% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/telemetry.md rename to python/docs/src/user-guide/core-user-guide/framework/telemetry.md index 1972cc4efeb2..6a7ede51a285 100644 --- a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/telemetry.md +++ b/python/docs/src/user-guide/core-user-guide/framework/telemetry.md @@ -3,9 +3,20 @@ AutoGen has native support for [open telemetry](https://opentelemetry.io/). This allows you to collect telemetry data from your application and send it to a telemetry backend of your choosing. These are the components that are currently instrumented: -- Runtime (Single Threaded Agent Runtime, Worker Agent Runtime) + +- Runtime ({py:class}`~autogen_core.SingleThreadedAgentRuntime` and {py:class}`~autogen_ext.runtimes.grpc.GrpcWorkerAgentRuntime`). +- Tool ({py:class}`~autogen_core.tools.BaseTool`) with the `execute_tool` span in [GenAI semantic convention for tools](https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span). +- AgentChat Agents ({py:class}`~autogen_agentchat.agents.BaseChatAgent`) with the `create_agent` and `invoke_agent` spans in [GenAI semantic convention for agents](https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-agent-spans/#create-agent-span). + +```{note} +To disable the agent runtime telemetry, you can set the `trace_provider` to +`opentelemetry.trace.NoOpTracerProvider` in the runtime constructor. + +Additionally, you can set the environment variable `AUTOGEN_DISABLE_RUNTIME_TRACING` to `true` to disable the agent runtime telemetry if you don't have access to the runtime constructor. For example, if you are using `ComponentConfig`. +``` ## Instrumenting your application + To instrument your application, you will need an sdk and an exporter. You may already have these if your application is already instrumented with open telemetry. ## Clean instrumentation @@ -26,6 +37,7 @@ pip install opentelemetry-exporter-otlp-proto-grpc ``` Next, we need to get a tracer provider: + ```python from opentelemetry import trace from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter @@ -44,6 +56,7 @@ def configure_oltp_tracing(endpoint: str = None) -> trace.TracerProvider: ``` Now you can send the trace_provider when creating your runtime: + ```python # for single threaded runtime single_threaded_runtime = SingleThreadedAgentRuntime(tracer_provider=tracer_provider) @@ -56,6 +69,7 @@ And that's it! Your application is now instrumented with open telemetry. You can ### Existing instrumentation If you have open telemetry already set up in your application, you can pass the tracer provider to the runtime when creating it: + ```python from opentelemetry import trace @@ -67,3 +81,8 @@ single_threaded_runtime = SingleThreadedAgentRuntime(tracer_provider=tracer_prov # or for worker runtime worker_runtime = GrpcWorkerAgentRuntime(tracer_provider=tracer_provider) ``` + +### Examples + +See [Tracing and Observability](../../agentchat-user-guide/tracing.ipynb) +for a complete example of how to set up open telemetry with AutoGen. diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/index.md b/python/docs/src/user-guide/core-user-guide/index.md similarity index 99% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/index.md rename to python/docs/src/user-guide/core-user-guide/index.md index ff51ffb355b1..dbbd4aaaac65 100644 --- a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/index.md +++ b/python/docs/src/user-guide/core-user-guide/index.md @@ -48,6 +48,7 @@ framework/component-config components/model-clients components/model-context components/tools +components/workbench components/command-line-code-executors ``` diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/installation.md b/python/docs/src/user-guide/core-user-guide/installation.md similarity index 94% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/installation.md rename to python/docs/src/user-guide/core-user-guide/installation.md index 641ba5f1ab8c..400b5e1ab0bd 100644 --- a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/installation.md +++ b/python/docs/src/user-guide/core-user-guide/installation.md @@ -10,11 +10,18 @@ When installing AgentChat locally, we recommend using a virtual environment for Create and activate: +Linux/Mac: ```bash python3 -m venv .venv source .venv/bin/activate ``` +Windows command-line: +```batch +python3 -m venv .venv +.venv\Scripts\activate.bat +``` + To deactivate later, run: ```bash diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/quickstart.ipynb b/python/docs/src/user-guide/core-user-guide/quickstart.ipynb similarity index 99% rename from python/packages/autogen-core/docs/src/user-guide/core-user-guide/quickstart.ipynb rename to python/docs/src/user-guide/core-user-guide/quickstart.ipynb index 29b696f31de6..6ee673e1beba 100644 --- a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/quickstart.ipynb +++ b/python/docs/src/user-guide/core-user-guide/quickstart.ipynb @@ -158,7 +158,7 @@ "source": [ "from autogen_core import AgentId, SingleThreadedAgentRuntime\n", "\n", - "# Create an local embedded runtime.\n", + "# Create a local embedded runtime.\n", "runtime = SingleThreadedAgentRuntime()\n", "\n", "# Register the modifier and checker agents by providing\n", diff --git a/python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/azure-container-code-executor.ipynb b/python/docs/src/user-guide/extensions-user-guide/azure-container-code-executor.ipynb similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/azure-container-code-executor.ipynb rename to python/docs/src/user-guide/extensions-user-guide/azure-container-code-executor.ipynb diff --git a/python/docs/src/user-guide/extensions-user-guide/azure-foundry-agent.ipynb b/python/docs/src/user-guide/extensions-user-guide/azure-foundry-agent.ipynb new file mode 100644 index 000000000000..c84360122579 --- /dev/null +++ b/python/docs/src/user-guide/extensions-user-guide/azure-foundry-agent.ipynb @@ -0,0 +1,131 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Azure AI Foundry Agent\n", + "\n", + "In AutoGen, you can build and deploy agents that are backed by the [Azure AI Foundry Agent Service](https://learn.microsoft.com/en-us/azure/ai-services/agents/overview) using the {py:class}`~autogen_ext.agents.azure._azure_ai_agent.AzureAIAgent` class. Here, important aspects of the agent including the provisioned model, tools (e.g, code interpreter, bing search grounding, file search etc.), observability, and security are managed by Azure. This allows you to focus on building your agent without worrying about the underlying infrastructure.\n", + "\n", + "In this guide, we will explore an example of creating an Azure AI Foundry Agent using the `AzureAIAgent` that can address tasks using the Azure Grounding with Bing Search tool." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# pip install \"autogen-ext[azure]\" # For Azure AI Foundry Agent Service" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Bing Search Grounding \n", + "\n", + "An {py:class}`~autogen_ext.agents.azure._azure_ai_agent.AzureAIAgent` can be assigned a set of tools including [Grounding with Bing Search](https://learn.microsoft.com/en-us/azure/ai-services/agents/how-to/tools/bing-grounding?tabs=python&pivots=overview#setup). \n", + "\n", + "Grounding with Bing Search allows your Azure AI Agents to incorporate real-time public web data when generating responses. You need to create a Grounding with Bing Search resource, and then connect this resource to your Azure AI Agents. When a user sends a query, Azure AI Agents decide if Grounding with Bing Search should be leveraged or not. If so, it will leverage Bing to search over public web data and return relevant chunks. Lastly, Azure AI Agents will use returned chunks to generate a response.\n", + "\n", + "## Prerequisites\n", + "\n", + "- You need to have an Azure subscription.\n", + "- You need to have the Azure CLI installed and configured. (also login using the command `az login` to enable default credentials)\n", + "- You need to have the `autogen-ext[azure]` package installed.\n", + "\n", + "You can create a [Grounding with Bing Search resource in the Azure portal](https://portal.azure.com/#create/Microsoft.BingGroundingSearch). Note that you will need to have owner or contributor role in your subscription or resource group to create it. Once you have created your resource, you can then pass it to the Azure Foundry Agent using the resource name.\n", + "\n", + "In the following example, we will create a new Azure Foundry Agent that uses the Grounding with Bing Search resource.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "import dotenv\n", + "from autogen_agentchat.messages import TextMessage\n", + "from autogen_core import CancellationToken\n", + "from autogen_ext.agents.azure import AzureAIAgent\n", + "from azure.ai.agents.models import BingGroundingTool\n", + "from azure.ai.projects.aio import AIProjectClient\n", + "from azure.identity.aio import DefaultAzureCredential\n", + "\n", + "dotenv.load_dotenv()\n", + "\n", + "\n", + "async def bing_example() -> None:\n", + " async with DefaultAzureCredential() as credential: # type: ignore\n", + " async with AIProjectClient( # type: ignore\n", + " credential=credential, endpoint=os.getenv(\"AZURE_PROJECT_ENDPOINT\", \"\")\n", + " ) as project_client:\n", + " conn = await project_client.connections.get(name=os.getenv(\"BING_CONNECTION_NAME\", \"\"))\n", + "\n", + " bing_tool = BingGroundingTool(conn.id)\n", + " agent_with_bing_grounding = AzureAIAgent(\n", + " name=\"bing_agent\",\n", + " description=\"An AI assistant with Bing grounding\",\n", + " project_client=project_client,\n", + " deployment_name=\"gpt-4o\",\n", + " instructions=\"You are a helpful assistant.\",\n", + " tools=bing_tool.definitions,\n", + " metadata={\"source\": \"AzureAIAgent\"},\n", + " )\n", + "\n", + " # For the bing grounding tool to return the citations, the message must contain an instruction for the model to do return them.\n", + " # For example: \"Please provide citations for the answers\"\n", + "\n", + " result = await agent_with_bing_grounding.on_messages(\n", + " messages=[\n", + " TextMessage(\n", + " content=\"What is Microsoft's annual leave policy? Provide citations for your answers.\",\n", + " source=\"user\",\n", + " )\n", + " ],\n", + " cancellation_token=CancellationToken(),\n", + " message_limit=5,\n", + " )\n", + " print(result)\n", + "\n", + "\n", + "await bing_example()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Note that you can also provide other Azure Backed [tools](https://learn.microsoft.com/en-us/azure/ai-services/agents/how-to/tools/overview) and local client side functions to the agent.\n", + "\n", + "See the {py:class}`~autogen_ext.agents.azure._azure_ai_agent.AzureAIAgent` class api documentation for more details on how to create an Azure Foundry Agent." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.12" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/create-your-own.md b/python/docs/src/user-guide/extensions-user-guide/create-your-own.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/create-your-own.md rename to python/docs/src/user-guide/extensions-user-guide/create-your-own.md diff --git a/python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/discover.md b/python/docs/src/user-guide/extensions-user-guide/discover.md similarity index 70% rename from python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/discover.md rename to python/docs/src/user-guide/extensions-user-guide/discover.md index fed1752a0da6..3040cc194ae3 100644 --- a/python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/discover.md +++ b/python/docs/src/user-guide/extensions-user-guide/discover.md @@ -42,6 +42,11 @@ Find community samples and examples of how to use AutoGen | [autogen-watsonx-client](https://github.com/tsinggggg/autogen-watsonx-client) | [PyPi](https://pypi.org/project/autogen-watsonx-client/) | Model client for [IBM watsonx.ai](https://www.ibm.com/products/watsonx-ai) | | [autogen-openaiext-client](https://github.com/vballoli/autogen-openaiext-client) | [PyPi](https://pypi.org/project/autogen-openaiext-client/) | Model client for other LLMs like Gemini, etc. through the OpenAI API | | [autogen-ext-mcp](https://github.com/richard-gyiko/autogen-ext-mcp) | [PyPi](https://pypi.org/project/autogen-ext-mcp/) | Tool adapter for Model Context Protocol server tools | +| [autogen-ext-email](https://github.com/masquerlin/autogen-ext-email) | [PyPi](https://pypi.org/project/autogen-ext-email/) | A Email agent for generating email and sending | +| [autogen-oaiapi](https://github.com/SongChiYoung/autogen-oaiapi) | [PyPi](https://pypi.org/project/autogen-oaiapi/) | an OpenAI-style API server built on top of AutoGen | +| [autogen-contextplus](https://github.com/SongChiYoung/autogen-contextplus) | [PyPi](https://pypi.org/project/autogen-contextplus/) | Enhanced model_context implementations, with features such as automatic summarization and truncation of model context. | +| [autogen-ext-yepcode](https://github.com/yepcode/autogen-ext-yepcode) | [PyPi](https://pypi.org/project/autogen-ext-yepcode/) | Enables agents to securely execute code in isolated remote sandboxes using [YepCode](https://yepcode.io)’s serverless runtime. | + diff --git a/python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/index.md b/python/docs/src/user-guide/extensions-user-guide/index.md similarity index 99% rename from python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/index.md rename to python/docs/src/user-guide/extensions-user-guide/index.md index c0f918ba2b19..964acbbc07d6 100644 --- a/python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/index.md +++ b/python/docs/src/user-guide/extensions-user-guide/index.md @@ -22,6 +22,7 @@ create-your-own :caption: Guides azure-container-code-executor +azure-foundry-agent ``` AutoGen is designed to be extensible. The `autogen-ext` package contains the built-in component implementations maintained by the AutoGen project. diff --git a/python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/installation.md b/python/docs/src/user-guide/extensions-user-guide/installation.md similarity index 100% rename from python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/installation.md rename to python/docs/src/user-guide/extensions-user-guide/installation.md diff --git a/python/packages/agbench/benchmarks/GAIA/Templates/ParallelAgents/expected_answer.txt b/python/packages/agbench/benchmarks/GAIA/Templates/ParallelAgents/expected_answer.txt new file mode 100644 index 000000000000..8153c2bf8242 --- /dev/null +++ b/python/packages/agbench/benchmarks/GAIA/Templates/ParallelAgents/expected_answer.txt @@ -0,0 +1 @@ +__EXPECTED_ANSWER__ diff --git a/python/packages/agbench/benchmarks/GAIA/Templates/ParallelAgents/prompt.txt b/python/packages/agbench/benchmarks/GAIA/Templates/ParallelAgents/prompt.txt new file mode 100644 index 000000000000..482f50dca311 --- /dev/null +++ b/python/packages/agbench/benchmarks/GAIA/Templates/ParallelAgents/prompt.txt @@ -0,0 +1 @@ +__PROMPT__ diff --git a/python/packages/agbench/benchmarks/GAIA/Templates/ParallelAgents/requirements.txt b/python/packages/agbench/benchmarks/GAIA/Templates/ParallelAgents/requirements.txt new file mode 100644 index 000000000000..3db8bfa55857 --- /dev/null +++ b/python/packages/agbench/benchmarks/GAIA/Templates/ParallelAgents/requirements.txt @@ -0,0 +1,5 @@ +tiktoken +pyyaml +/autogen_python/packages/autogen-core +/autogen_python/packages/autogen-ext[openai,magentic-one] +/autogen_python/packages/autogen-agentchat diff --git a/python/packages/agbench/benchmarks/GAIA/Templates/ParallelAgents/scenario.py b/python/packages/agbench/benchmarks/GAIA/Templates/ParallelAgents/scenario.py new file mode 100644 index 000000000000..9922b33bc1a1 --- /dev/null +++ b/python/packages/agbench/benchmarks/GAIA/Templates/ParallelAgents/scenario.py @@ -0,0 +1,402 @@ +import asyncio +import os +import re +import logging +import yaml +import warnings +import contextvars +import builtins +import shutil +import json +from datetime import datetime +from typing import List, Optional, Dict +from collections import deque +from autogen_agentchat import TRACE_LOGGER_NAME as AGENTCHAT_TRACE_LOGGER_NAME, EVENT_LOGGER_NAME as AGENTCHAT_EVENT_LOGGER_NAME +from autogen_core import TRACE_LOGGER_NAME as CORE_TRACE_LOGGER_NAME, EVENT_LOGGER_NAME as CORE_EVENT_LOGGER_NAME +from autogen_ext.agents.magentic_one import MagenticOneCoderAgent +from autogen_agentchat.teams import MagenticOneGroupChat +from autogen_agentchat.ui import Console +from autogen_core.models import ( + AssistantMessage, + ChatCompletionClient, + LLMMessage, + UserMessage, +) +from autogen_core.logging import LLMCallEvent +from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor +from autogen_agentchat.conditions import TextMentionTermination +from autogen_core.models import ChatCompletionClient +from autogen_ext.agents.web_surfer import MultimodalWebSurfer +from autogen_ext.agents.file_surfer import FileSurfer +from autogen_agentchat.agents import CodeExecutorAgent +from autogen_agentchat.messages import ( + TextMessage, + AgentEvent, + ChatMessage, + HandoffMessage, + MultiModalMessage, + StopMessage, + TextMessage, + ToolCallExecutionEvent, + ToolCallRequestEvent, + ToolCallSummaryMessage, +) +from autogen_core import CancellationToken +from autogen_ext.models.openai import OpenAIChatCompletionClient +from autogen_ext.models.openai._model_info import _MODEL_TOKEN_LIMITS, resolve_model +from autogen_agentchat.utils import content_to_str + +# Suppress warnings about the requests.Session() not being closed +warnings.filterwarnings(action="ignore", message="unclosed", category=ResourceWarning) + +core_event_logger = logging.getLogger(CORE_EVENT_LOGGER_NAME) +agentchat_event_logger = logging.getLogger(AGENTCHAT_EVENT_LOGGER_NAME) +agentchat_trace_logger = logging.getLogger(AGENTCHAT_TRACE_LOGGER_NAME) + +# Create a context variable to hold the current team's log file and the current team id. +current_log_file = contextvars.ContextVar("current_log_file", default=None) +current_team_id = contextvars.ContextVar("current_team_id", default=None) + +# Save the original print function and event_logger.info method. +original_print = builtins.print +original_agentchat_event_logger_info = agentchat_event_logger.info +original_core_event_logger_info = core_event_logger.info + +class LogHandler(logging.FileHandler): + def __init__(self, filename: str = "log.jsonl", print_message: bool = True) -> None: + super().__init__(filename, mode="w") + self.print_message = print_message + + def emit(self, record: logging.LogRecord) -> None: + try: + ts = datetime.fromtimestamp(record.created).isoformat() + if AGENTCHAT_EVENT_LOGGER_NAME in record.name: + original_msg = record.msg + record.msg = json.dumps( + { + "timestamp": ts, + "source": record.msg.source, + "message": content_to_str(record.msg.content), + "type": record.msg.type, + } + ) + super().emit(record) + record.msg = original_msg + elif CORE_EVENT_LOGGER_NAME in record.name: + if isinstance(record.msg, LLMCallEvent): + original_msg = record.msg + record.msg = json.dumps( + { + "timestamp": ts, + "prompt_tokens": record.msg.kwargs["prompt_tokens"], + "completion_tokens": record.msg.kwargs["completion_tokens"], + "type": "LLMCallEvent", + } + ) + super().emit(record) + record.msg = original_msg + except Exception: + print("error in logHandler.emit", flush=True) + self.handleError(record) + +def tee_print(*args, **kwargs): + # Get the current log file from the context. + log_file = current_log_file.get() + # Call the original print (goes to the console). + original_print(*args, **kwargs) + # Also write to the log file if one is set. + if log_file is not None: + sep = kwargs.get("sep", " ") + end = kwargs.get("end", "\n") + message = sep.join(map(str, args)) + end + log_file.write(message) + log_file.flush() + +def team_specific_agentchat_event_logger_info(msg, *args, **kwargs): + team_id = current_team_id.get() + if team_id is not None: + # Get a logger with a team-specific name. + team_logger = logging.getLogger(f"{AGENTCHAT_EVENT_LOGGER_NAME}.team{team_id}") + team_logger.info(msg, *args, **kwargs) + else: + original_agentchat_event_logger_info(msg, *args, **kwargs) + +def team_specific_core_event_logger_info(msg, *args, **kwargs): + team_id = current_team_id.get() + if team_id is not None: + # Get a logger with a team-specific name. + team_logger = logging.getLogger(f"{CORE_EVENT_LOGGER_NAME}.team{team_id}") + team_logger.info(msg, *args, **kwargs) + else: + original_core_event_logger_info(msg, *args, **kwargs) + +# Monkey-patch the built-in print and event_logger.info methods with our team-specific versions. +builtins.print = tee_print +agentchat_event_logger.info = team_specific_agentchat_event_logger_info +core_event_logger.info = team_specific_core_event_logger_info + +async def run_team(team: MagenticOneGroupChat, team_idx: int, task: str, cancellation_token: CancellationToken, logfile): + token_logfile = current_log_file.set(logfile) + token_team_id = current_team_id.set(team_idx) + try: + task_result = await Console( + team.run_stream( + task=task.strip(), + cancellation_token=cancellation_token + ) + ) + return team_idx, task_result + finally: + current_log_file.reset(token_logfile) + current_team_id.reset(token_team_id) + logfile.close() + +async def aggregate_final_answer(task: str, client: ChatCompletionClient, team_results, source: str = "Aggregator", cancellation_token: Optional[CancellationToken] = None) -> str: + """ + team_results: {"team_key": TaskResult} + team_completion_order: The order in which the teams completed their tasks + """ + + if len(team_results) == 1: + final_answer = list(team_results.values())[0].messages[-1].content + aggregator_logger.info( + f"{source} (Response):\n{final_answer}" + ) + return final_answer + + assert len(team_results) > 1 + + aggregator_messages_to_send = {team_id: deque() for team_id in team_results.keys()} # {team_id: context} + + team_ids = list(team_results.keys()) + current_round = 0 + while ( + not all(len(team_result.messages) == 0 for team_result in team_results.values()) + and ((not resolve_model(client._create_args["model"]) in _MODEL_TOKEN_LIMITS) or client.remaining_tokens([m for messages in aggregator_messages_to_send.values() for m in messages]) + > 2000) + ): + team_idx = team_ids[current_round % len(team_ids)] + if len(team_results[team_idx].messages) > 0: + m = team_results[team_idx].messages[-1] + if isinstance(m, ToolCallRequestEvent | ToolCallExecutionEvent): + # Ignore tool call messages. + pass + elif isinstance(m, StopMessage | HandoffMessage): + aggregator_messages_to_send[team_idx].appendleft(UserMessage(content=m.to_model_text(), source=m.source)) + elif m.source == "MagenticOneOrchestrator": + assert isinstance(m, TextMessage | ToolCallSummaryMessage) + aggregator_messages_to_send[team_idx].appendleft(AssistantMessage(content=m.to_model_text(), source=m.source)) + else: + assert isinstance(m, (TextMessage, MultiModalMessage, ToolCallSummaryMessage)) + aggregator_messages_to_send[team_idx].appendleft(UserMessage(content=m.to_model_text(), source=m.source)) + team_results[team_idx].messages.pop() + current_round += 1 + + # Log the messages to send + payload = "" + for team_idx, messages in aggregator_messages_to_send.items(): + payload += f"\n{'*'*75} \n" f"Team #: {team_idx}" f"\n{'*'*75} \n" + for message in messages: + payload += f"\n{'-'*75} \n" f"{message.source}:\n" f"\n{message.content}\n" + payload += f"\n{'-'*75} \n" f"Team #{team_idx} stop reason:\n" f"\n{team_results[team_idx].stop_reason}\n" + payload += f"\n{'*'*75} \n" + aggregator_logger.info(f"{source} (Aggregator Messages):\n{payload}") + + context: List[LLMMessage] = [] + + # Add the preamble + context.append( + UserMessage( + content=f"Earlier you were asked the following:\n\n{task}\n\nYour team then worked diligently to address that request. You have been provided with a collection of transcripts and stop reasons from {len(team_results)} different teams to the question. Your task is to carefully evaluate the correctness of each team's response by analyzing their respective transcripts and stop reasons. After considering all perspectives, provide a FINAL ANSWER to the question. It is crucial to critically evaluate the information provided in these responses, recognizing that some of it may be biased or incorrect.", + source=source, + ) + ) + + for team_idx, aggregator_messages in aggregator_messages_to_send.items(): + context.append( + UserMessage( + content=f"Transcript from Team #{team_idx}:", + source=source, + ) + ) + for message in aggregator_messages: + context.append(message) + context.append( + UserMessage( + content=f"Stop reason from Team #{team_idx}:", + source=source, + ) + ) + context.append( + UserMessage( + content=team_results[team_idx].stop_reason if team_results[team_idx].stop_reason else "No stop reason provided.", + source=source, + ) + ) + + # ask for the final answer + context.append( + UserMessage( + content=f""" + Let's think step-by-step. Carefully review the conversation above, critically evaluate the correctness of each team's response, and then output a FINAL ANSWER to the question. The question is repeated here for convenience: + + {task} + + To output the final answer, use the following template: FINAL ANSWER: [YOUR FINAL ANSWER] + Your FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. + ADDITIONALLY, your FINAL ANSWER MUST adhere to any formatting instructions specified in the original question (e.g., alphabetization, sequencing, units, rounding, decimal places, etc.) + If you are asked for a number, express it numerically (i.e., with digits rather than words), don't use commas, and don't include units such as $ or percent signs unless specified otherwise. + If you are asked for a string, don't use articles or abbreviations (e.g. for cities), unless specified otherwise. Don't output any final sentence punctuation such as '.', '!', or '?'. + If you are asked for a comma separated list, apply the above rules depending on whether the elements are numbers or strings. + """.strip(), + source=source, + ) + ) + + response = await client.create(context, cancellation_token=cancellation_token) + assert isinstance(response.content, str) + + final_answer = re.sub(r"FINAL ANSWER:", "[FINAL ANSWER]:", response.content) + aggregator_logger.info( + f"{source} (Response):\n{final_answer}" + ) + + return re.sub(r"FINAL ANSWER:", "FINAL AGGREGATED ANSWER:", response.content) + + +async def main(num_teams: int, num_answers: int) -> None: + + # Load model configuration and create the model client. + with open("config.yaml", "r") as f: + config = yaml.safe_load(f) + + orchestrator_client = ChatCompletionClient.load_component(config["orchestrator_client"]) + coder_client = ChatCompletionClient.load_component(config["coder_client"]) + web_surfer_client = ChatCompletionClient.load_component(config["web_surfer_client"]) + file_surfer_client = ChatCompletionClient.load_component(config["file_surfer_client"]) + + # Read the prompt + prompt = "" + with open("prompt.txt", "rt") as fh: + prompt = fh.read().strip() + filename = "__FILE_NAME__".strip() + + # Prepare the prompt + filename_prompt = "" + if len(filename) > 0: + filename_prompt = f"The question is about a file, document or image, which can be accessed by the filename '{filename}' in the current working directory." + task = f"{prompt}\n\n{filename_prompt}" + + # Reset logs directory (remove all files in it) + logs_dir = "logs" + if os.path.exists(logs_dir): + shutil.rmtree(logs_dir) + + teams = [] + async_tasks = [] + tokens = [] + for team_idx in range(num_teams): + # Set up the team + coder = MagenticOneCoderAgent( + "Assistant", + model_client = coder_client, + ) + + executor = CodeExecutorAgent("ComputerTerminal", code_executor=LocalCommandLineCodeExecutor()) + + file_surfer = FileSurfer( + name="FileSurfer", + model_client = file_surfer_client, + ) + + web_surfer = MultimodalWebSurfer( + name="WebSurfer", + model_client = web_surfer_client, + downloads_folder=os.getcwd(), + debug_dir=logs_dir, + to_save_screenshots=True, + ) + team = MagenticOneGroupChat( + [coder, executor, file_surfer, web_surfer], + model_client=orchestrator_client, + max_turns=30, + final_answer_prompt= f""", +We have completed the following task: + +{prompt} + +The above messages contain the conversation that took place to complete the task. +Read the above conversation and output a FINAL ANSWER to the question. +To output the final answer, use the following template: FINAL ANSWER: [YOUR FINAL ANSWER] +Your FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. +ADDITIONALLY, your FINAL ANSWER MUST adhere to any formatting instructions specified in the original question (e.g., alphabetization, sequencing, units, rounding, decimal places, etc.) +If you are asked for a number, express it numerically (i.e., with digits rather than words), don't use commas, and don't include units such as $ or percent signs unless specified otherwise. +If you are asked for a string, don't use articles or abbreviations (e.g. for cities), unless specified otherwise. Don't output any final sentence punctuation such as '.', '!', or '?'. +If you are asked for a comma separated list, apply the above rules depending on whether the elements are numbers or strings. +""".strip() + ) + teams.append(team) + cancellation_token = CancellationToken() + tokens.append(cancellation_token) + logfile = open(f"console_log_{team_idx}.txt", "w") + team_agentchat_logger = logging.getLogger(f"{AGENTCHAT_EVENT_LOGGER_NAME}.team{team_idx}") + team_core_logger = logging.getLogger(f"{CORE_EVENT_LOGGER_NAME}.team{team_idx}") + team_log_handler = LogHandler(f"log_{team_idx}.jsonl", print_message=False) + team_agentchat_logger.addHandler(team_log_handler) + team_core_logger.addHandler(team_log_handler) + async_task = asyncio.create_task( + run_team(team, team_idx, task, cancellation_token, logfile) + ) + async_tasks.append(async_task) + + # Wait until at least num_answers tasks have completed. + team_results = {} + for future in asyncio.as_completed(async_tasks): + try: + team_id, result = await future + team_results[team_id] = result + except Exception as e: + # Optionally log exception. + print(f"Task raised an exception: {e}") + if len(team_results) >= num_answers: + break + + # Cancel any pending teams. + for task, token in zip(async_tasks, tokens): + if not task.done(): + token.cancel() + # Await all tasks to handle cancellation gracefully. + await asyncio.gather(*async_tasks, return_exceptions=True) + + print("len(team_results):", len(team_results)) + final_answer = await aggregate_final_answer(prompt, orchestrator_client, team_results) + print(final_answer) + +if __name__ == "__main__": + num_teams = 3 + num_answers = 3 + + agentchat_trace_logger.setLevel(logging.DEBUG) + file_handler = logging.FileHandler("trace.log", mode="w") + file_handler.setLevel(logging.DEBUG) + formatter = logging.Formatter( + "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + ) + file_handler.setFormatter(formatter) + agentchat_trace_logger.addHandler(file_handler) + + core_event_logger.setLevel(logging.DEBUG) + agentchat_event_logger.setLevel(logging.DEBUG) + log_handler = LogHandler() + core_event_logger.addHandler(log_handler) + agentchat_event_logger.addHandler(log_handler) + + # Create another logger for the aggregator + aggregator_logger = logging.getLogger("aggregator") + aggregator_logger.setLevel(logging.DEBUG) + fh = logging.FileHandler("aggregator_log.txt", mode="w") + fh.setLevel(logging.DEBUG) + aggregator_logger.addHandler(fh) + + + asyncio.run(main(num_teams, num_answers)) diff --git a/python/packages/agbench/benchmarks/GAIA/Templates/SelectorGroupChat/scenario.py b/python/packages/agbench/benchmarks/GAIA/Templates/SelectorGroupChat/scenario.py index e2e1d8fae009..5fa4b00273f2 100644 --- a/python/packages/agbench/benchmarks/GAIA/Templates/SelectorGroupChat/scenario.py +++ b/python/packages/agbench/benchmarks/GAIA/Templates/SelectorGroupChat/scenario.py @@ -16,7 +16,7 @@ from autogen_ext.agents.web_surfer import MultimodalWebSurfer from autogen_ext.agents.file_surfer import FileSurfer from autogen_agentchat.agents import CodeExecutorAgent -from autogen_agentchat.messages import TextMessage, AgentEvent, ChatMessage, HandoffMessage, MultiModalMessage, StopMessage +from autogen_agentchat.messages import TextMessage, BaseAgentEvent, BaseChatMessage, HandoffMessage, MultiModalMessage, StopMessage from autogen_core.models import LLMMessage, UserMessage, AssistantMessage # Suppress warnings about the requests.Session() not being closed @@ -141,7 +141,7 @@ def __init__(self, prompt: str, model_client: ChatCompletionClient, termination_ def terminated(self) -> bool: return self._terminated - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: if self._terminated: raise TerminatedException("Termination condition has already been reached") diff --git a/python/packages/agbench/src/agbench/linter/__init__.py b/python/packages/agbench/src/agbench/linter/__init__.py index 797b7f272a5b..a104962445f6 100644 --- a/python/packages/agbench/src/agbench/linter/__init__.py +++ b/python/packages/agbench/src/agbench/linter/__init__.py @@ -1,4 +1,4 @@ # __init__.py -from ._base import Code, Document, CodedDocument, BaseQualitativeCoder +from ._base import BaseQualitativeCoder, Code, CodedDocument, Document __all__ = ["Code", "Document", "CodedDocument", "BaseQualitativeCoder"] diff --git a/python/packages/agbench/src/agbench/linter/_base.py b/python/packages/agbench/src/agbench/linter/_base.py index 4f6209b7809c..ea0a893e4e7a 100644 --- a/python/packages/agbench/src/agbench/linter/_base.py +++ b/python/packages/agbench/src/agbench/linter/_base.py @@ -1,12 +1,14 @@ -import json import hashlib +import json import re -from typing import Protocol, List, Set, Optional +from typing import List, Optional, Protocol, Set + from pydantic import BaseModel, Field class Document(BaseModel): text: str = Field(..., description="Text content of the document.") + lines: List[str] = Field(..., description="List of lines in the document. This is a list of strings.") name: Optional[str] = Field(None, description="Optional name of the document.") def __hash__(self) -> int: diff --git a/python/packages/agbench/src/agbench/linter/cli.py b/python/packages/agbench/src/agbench/linter/cli.py index 426890258b69..2a2c28aa9bab 100644 --- a/python/packages/agbench/src/agbench/linter/cli.py +++ b/python/packages/agbench/src/agbench/linter/cli.py @@ -1,8 +1,10 @@ -import os import argparse -from typing import List, Sequence, Optional +import os +from typing import List, Optional, Sequence + from openai import OpenAI -from ._base import Document, CodedDocument + +from ._base import CodedDocument, Document from .coders.oai_coder import OAIQualitativeCoder @@ -23,7 +25,7 @@ def load_log_file(path: str, prepend_numbers: bool = False) -> Document: lines = prepend_line_numbers(lines) text = "".join(lines) - return Document(text=text, name=os.path.abspath(path)) + return Document(text=text, lines=lines, name=os.path.abspath(path)) def code_log(path: str) -> Optional[CodedDocument]: diff --git a/python/packages/agbench/src/agbench/linter/coders/_prompt.py b/python/packages/agbench/src/agbench/linter/coders/_prompt.py new file mode 100644 index 000000000000..7402c4100a99 --- /dev/null +++ b/python/packages/agbench/src/agbench/linter/coders/_prompt.py @@ -0,0 +1,73 @@ +MAIN_PROMPT = """You are an expert qualitative researcher. + +Given a document containing errors below, generate a list of (error) codes. +The document shows a log of interaction between multiple agents collaborating +to solve a complex task. + +For example, the name could be of the format "lack-of-word2", +"failed-to-bar", "excessive-use-of-magenta". Name should adhere to +Joseph M. Williams' writing principles of clarity, conciseness, and coherence. + +Ensure each code name is lower-case, hyphenated, and directly reflects the +concept it represents. Avoid ambiguous or overly complex terms, and prioritize +simplicity, precision, and readability in the naming. + +The code names should pass the 'clarity and grace' test by being easy to +understand, descriptive, and reflective of the content they categorize. +- suggest codes that are similar to good code names. avoid code names that are +similar to bad code names. +- The definition should be simple worded and practical. At least 2 sentences, +max 3. It should be written in past tense. + +It should convey how a labeller could apply this code to future logs, without +mentioning the word "labeller". The definition should be specific enough to be +useful in debugging. It should be very concrete. And should be well thought and +make sense. Bull shitting will not earn you any points. + +- The examples should be a list. Each example should be descriptive between +2-3 sentences. Examples should be concrete, informative and not vague. Provide +at max 20 salient examples. Examples should contain a lot of detail about what +happened and should refer to incidents in the log. + +- The list of codes must mutually exclusive. + +# GOOD EXAMPLES OF FINAL CODE NAMES/CLUSTERS +* looped-without-progress +* repeated-unsuccessful-actions +* repeated-syntax-errors +* exceeded-context-window-limits +* encountered-security-risks +* failure-to-switch-strategy +* exceeded-resource-limits +* attempted-to-handle-excessive-data +* no-errors-detected +These names are high-level but also concrete. They exactly mention the type of +error, issue, gap that has been identified. + +## BAD EXAMPLES OF FINAL CODE NAMES/CLUSTERS +* mismanaged-data-utilization -- too high level +* incomplete-or-misguided-execution -- too high level +* misaligned-agent-interactions -- too high level +* mismanaged-task-strategies -- too high level +* resource-inefficiencies -- vague +* communication-issues -- vague +* coordination-issues -- too high level and vague +* operational-failures +* execution-errors -- too high level +* navigation-issues -- too concise +* adaptive-failures -- too concise +* successful-processes -- I dont like the word processes +* system-constraints +* configuration-issues +* information-inaccuracies -- too high level +* process-improvements -- vague, not an error +* inadequate-error-response -- too high-level, unclear what kind of errors +* specific-access-issues -- makes no sense +* strategy-inefficiency -- strategy is too high level +* error-management-gaps -- unclear what error management means +* error-handling-deficiency -- unclear what kind of errors +* coordination-breakdown -- unclear what coordination means +* muddled-task-execution -- unclear what kind of tasks were muddled +* task-completion-gaps -- too high level +The above names are too high level and unclear. Please DO NOT use such names. +""" diff --git a/python/packages/agbench/src/agbench/linter/coders/oai_coder.py b/python/packages/agbench/src/agbench/linter/coders/oai_coder.py index 374093d3d81b..451abdfed4e2 100644 --- a/python/packages/agbench/src/agbench/linter/coders/oai_coder.py +++ b/python/packages/agbench/src/agbench/linter/coders/oai_coder.py @@ -1,13 +1,12 @@ import os import re - -from typing import List, Set, Optional -from pydantic import BaseModel +from typing import List, Optional, Set from openai import OpenAI +from pydantic import BaseModel -from .._base import CodedDocument, Document, Code -from .._base import BaseQualitativeCoder +from .._base import BaseQualitativeCoder, Code, CodedDocument, CodeExample, Document +from ._prompt import MAIN_PROMPT class CodeList(BaseModel): @@ -23,6 +22,7 @@ def remove_control_characters(text: str) -> str: class OAIQualitativeCoder(BaseQualitativeCoder): DEFAULT_MODEL = "gpt-4o" + MAIN_PROMPT = MAIN_PROMPT def __init__(self, cache_dir: str = ".cache", model: str = DEFAULT_MODEL, cache_enabled: bool = False) -> None: self.client = OpenAI() @@ -30,11 +30,72 @@ def __init__(self, cache_dir: str = ".cache", model: str = DEFAULT_MODEL, cache_ self.model = model self.cache_enabled = cache_enabled - def code_document( - self, - doc: Document, - code_set: Optional[Set[Code]] = None, - ) -> Optional[CodedDocument]: + def code_document(self, doc: Document, code_set: Optional[Set[Code]] = None) -> Optional[CodedDocument]: + coded_doc = self._code_document(doc) + if coded_doc is None: + raise ValueError("Error in coding document with OpenAI") + + feedback = self._reflect_on_codes(coded_doc) + + coded_doc = self._code_document_with_feedback(coded_doc, feedback) + + if coded_doc is None: + raise ValueError("Error in coding document with OpenAI") + + feedback = self._reflect_on_codes(coded_doc) + + coded_doc = self._code_document_with_feedback(coded_doc, feedback) + + return coded_doc + + def _code_document_with_feedback(self, coded_doc: CodedDocument, feedback: str) -> Optional[CodedDocument]: + """ + Given a coded document and feedback, update the codes in the document. + + Again uses completion to generate new code lists + based on the doc, original codes, and feedback. + """ + + prompt = self.MAIN_PROMPT + + prompt += "\nDocument:\n" + for line in coded_doc.doc.lines: + prompt += f"{line}" + prompt += "Notice that the document contains the following number of lines: " + prompt += str(len(coded_doc.doc.lines)) + + prompt += "\n\n" + + prompt += "A previous attempt to code the document resulted in the following codes:\n" + for code in coded_doc.codes: + prompt += code.model_dump_json(indent=4) + prompt += "\n" + prompt += "\n\n" + + prompt += "A human expert has provided the following feedback on the codes:\n" + prompt += f"{feedback}\n\n" + + prompt += "Now revise the codes based on the feedback. " + + # save coding with feedback prompt to a file + # with open("coding_with_feedback_prompt.txt", "w") as f: + # f.write(prompt) + + completion = self.client.beta.chat.completions.parse( + model=self.model, + messages=[{"role": "user", "content": prompt}], + response_format=CodeList, + ) + message = completion.choices[0].message + if message.parsed and len(message.parsed.code_list) > 0: + coded_doc.codes = set(message.parsed.code_list) + else: + print(message.refusal) + raise ValueError("Error in coding document with OpenAI") + + return coded_doc + + def _code_document(self, doc: Document) -> Optional[CodedDocument]: # get hash of the document doc_hash = hash(doc) cache_file = os.path.join(self.cache_dir, f"{doc_hash}.json") if self.cache_enabled else None @@ -52,15 +113,12 @@ def code_document( coded_document: Optional[CodedDocument] = None - if code_set is None: - completion = self.client.beta.chat.completions.parse( - model=self.model, - messages=[ - { - "role": "system", - "content": """You are an expert qualitative researcher. + prompt = """You are an expert qualitative researcher. + +Given a document containing errors below, generate a list of (error) codes. +The document shows a log of interaction between multiple agents collaborating +to solve a complex task. -Given a list of dcocuments containing errors below, generate a list of (error) codes. Each code should contains: - at least 3 words, max 4 word, hyphenated. @@ -77,7 +135,7 @@ def code_document( - suggest codes that are similar to good code names. avoid code names that are similar to bad code names. - The definition should be simple worded and practical. At least 2 sentences, - max 3. It should be written in past tense. +max 3. It should be written in past tense. It should convey how a labeller could apply this code to future logs, without mentioning the word "labeller". The definition should be specific enough to be @@ -130,81 +188,134 @@ def code_document( * muddled-task-execution -- unclear what kind of tasks were muddled * task-completion-gaps -- too high level The above names are too high level and unclear. Please DO NOT use such names. - """, - }, - { - "role": "user", - "content": doc.text, - }, - ], - response_format=CodeList, - ) - - message = completion.choices[0].message - if message.parsed and len(message.parsed.code_list) > 0: - coded_document = CodedDocument(doc=doc, codes=set(message.parsed.code_list)) - else: - print(message.refusal) - raise ValueError("Error in coding document with OpenAI") + +Document: + +""" + + for line in doc.lines: + prompt += f"{line}" + prompt += "\n\n" + prompt += "Notice that the document contains the following number of lines: " + prompt += str(len(doc.lines)) + prompt += "\n\n" + + prompt += ( + "Now generate a list of codes for the document." + " Especially codes that detect errors/inefficiencies in the document." + ) + + # save the coding prompt to a file + # with open("coding_prompt.txt", "w") as f: + # f.write(prompt) + + completion = self.client.beta.chat.completions.parse( + model=self.model, + messages=[ + { + "role": "user", + "content": prompt, + }, + ], + response_format=CodeList, + ) + + message = completion.choices[0].message + if message.parsed and len(message.parsed.code_list) > 0: + coded_document = CodedDocument(doc=doc, codes=set(message.parsed.code_list)) else: - code_to_str = "\n".join( - [ - ( - f"\n---\nCode Name: {code.name}\n" - f"Definition: {code.definition}\n" - f"Examples: {code.examples}\n---\n" - ) - for code in code_set - ] - ) - - completion = self.client.beta.chat.completions.parse( - model=self.model, - messages=[ - { - "role": "system", - "content": """You are an expert qualitative researcher. - You can answer any questions about coding logs.""", - }, - { - "role": "user", - "content": f""" -## Context -The text below shows a log containing errors. Your task is to code the log with -the following codes. Generate a list of codes for the log below. - -Only use the codes from the list below. Do not create new codes. -Modify the examples of the codes to fit the context of the log. - -Your example should be informative to narrow down the details of the error in -the context of the example. - -## Codes - -{code_to_str} - -## Log - -{doc.text} -""", - }, - ], - response_format=CodeList, - ) - - message = completion.choices[0].message - if message.parsed and len(message.parsed.code_list) > 0: - code_list = message.parsed.code_list - # filter out codes whose names are not in the code_set - code_set_names = {code.name for code in code_set} - code_list = [code for code in code_list if code.name in code_set_names] - - coded_document = CodedDocument(doc=doc, codes=set(code_list)) - - if coded_document is None: + print(message.refusal) raise ValueError("Error in coding document with OpenAI") if self.cache_enabled and cache_file: with open(cache_file, "w") as f: f.write(coded_document.model_dump_json(indent=4)) + return coded_document + + def _codes_to_string(self, codes: Set[Code]) -> str: + """ + Convert a set of codes to a string representation. + Include name, definition, examples, line number, and severity. + """ + code_list: List[str] = [] + for code in codes: + code_list.append(f"[{code.severity}]: {code.name}: {code.definition}") + for example in code.examples: + code_list.append(f"\t{example.line}:{example.line_end}\t{example.reason}") + return "\n".join(code_list) + + def _extract_lines(self, doc: Document, start: int, end: int, buffer: int = 1) -> str: + """ + Extract a line from the document. + """ + start_line = max(0, start - buffer) + end_line = min(len(doc.lines), end + buffer) + lines = doc.lines[start_line:end_line] + return "".join(lines) + + def _extract_code_lines(self, doc: Document, example: CodeExample) -> str: + """ + Extract lines from the document based on the code. + """ + start = example.line + end = example.line_end + lines = self._extract_lines(doc, start, end) + return lines + + def _reflect_on_codes(self, coded_doc: CodedDocument) -> str: + """ + Given a coded document generate feedback. + E.g., whether the code used seem appropriate or not. + """ + + prompt = ( + "You are an expert qualitative researcher. " + "You are given a list of codes. " + "Pay attention the codes and the lines mentioned in the examples of the codes. " + "Which examples fail to spot meaningful errors? " + "Be direct and critical. " + "If a code identifies a BS error, say it. " + "There is no need to figure out how to fix the actual error. " + "The goal is to double check the validity of detected errors.\n\n" + ) + + # for line in coded_doc.doc.lines: + # prompt += f"{line}" + # prompt += "\n\n" + + # prompt += "Notice that the document contains the following number of lines: " + # prompt += str(len(coded_doc.doc.lines)) + + # prompt += "\n\n" + prompt += "A qualitative coding of a document claims to spot the following errors:\n\n" + for code in coded_doc.codes: + prompt += f"Code: {code.name}\n" + prompt += f"Definition: {code.definition}\n" + prompt += "Examples:\n" + for example in code.examples: + extracted_lines = self._extract_code_lines(coded_doc.doc, example) + prompt += f"- Does the text in the lines {example.line}:{example.line_end} shown below have enough information to justify the {code.name} error? " + prompt += f"Especially does the line {example.line} contain the error?\n\n" + prompt += f"{extracted_lines}\n" + prompt += "\n\n" + prompt += "\n" + + prompt += ( + "Now carefully analyze the examples. And provide feedback on the codes." + "If the examples lines do not align with the code name or definition, provide feedback." + ) + + # save the reflection_prompt to a file + # with open("reflection_prompt.txt", "w") as f: + # f.write(prompt) + + completion = self.client.chat.completions.create( + model=self.model, + messages=[{"role": "user", "content": prompt}], + ) + + feedback = completion.choices[0].message.content + if feedback is None: + raise ValueError("Error in generating feedback with OpenAI") + return feedback diff --git a/python/packages/autogen-agentchat/pyproject.toml b/python/packages/autogen-agentchat/pyproject.toml index 9a1b773843d1..978989805deb 100644 --- a/python/packages/autogen-agentchat/pyproject.toml +++ b/python/packages/autogen-agentchat/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "autogen-agentchat" -version = "0.4.9" +version = "0.7.2" license = {file = "LICENSE-CODE"} description = "AutoGen agents and teams library" readme = "README.md" @@ -15,7 +15,7 @@ classifiers = [ "Operating System :: OS Independent", ] dependencies = [ - "autogen-core==0.4.9", + "autogen-core==0.7.2", ] [tool.ruff] diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/__init__.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/__init__.py index a6732d12ef3f..ebce7b8e3baf 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/__init__.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/__init__.py @@ -5,7 +5,8 @@ from ._assistant_agent import AssistantAgent from ._base_chat_agent import BaseChatAgent -from ._code_executor_agent import CodeExecutorAgent +from ._code_executor_agent import ApprovalFuncType, ApprovalRequest, ApprovalResponse, CodeExecutorAgent +from ._message_filter_agent import MessageFilterAgent, MessageFilterConfig, PerSourceFilter from ._society_of_mind_agent import SocietyOfMindAgent from ._user_proxy_agent import UserProxyAgent @@ -15,4 +16,10 @@ "CodeExecutorAgent", "SocietyOfMindAgent", "UserProxyAgent", + "MessageFilterAgent", + "MessageFilterConfig", + "PerSourceFilter", + "ApprovalRequest", + "ApprovalResponse", + "ApprovalFuncType", ] diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py index e7ff6cc291a3..b1b1c7b57d31 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py @@ -1,6 +1,9 @@ +from __future__ import annotations + import asyncio import json import logging +import uuid import warnings from typing import ( Any, @@ -13,6 +16,7 @@ Optional, Sequence, Tuple, + TypeVar, Union, ) @@ -29,23 +33,23 @@ FunctionExecutionResult, FunctionExecutionResultMessage, LLMMessage, - ModelFamily, SystemMessage, - UserMessage, ) -from autogen_core.tools import BaseTool, FunctionTool -from pydantic import BaseModel +from autogen_core.tools import BaseTool, FunctionTool, StaticStreamWorkbench, ToolResult, Workbench +from pydantic import BaseModel, Field from typing_extensions import Self from .. import EVENT_LOGGER_NAME from ..base import Handoff as HandoffBase from ..base import Response from ..messages import ( - AgentEvent, - ChatMessage, + BaseAgentEvent, + BaseChatMessage, HandoffMessage, MemoryQueryEvent, ModelClientStreamingChunkEvent, + StructuredMessage, + StructuredMessageFactory, TextMessage, ThoughtEvent, ToolCallExecutionEvent, @@ -58,13 +62,18 @@ event_logger = logging.getLogger(EVENT_LOGGER_NAME) +# Add type variables for more specific typing +T = TypeVar("T", bound=BaseModel) +R = TypeVar("R", bound=BaseModel) + class AssistantAgentConfig(BaseModel): """The declarative configuration for the assistant agent.""" name: str model_client: ComponentModel - tools: List[ComponentModel] | None + tools: List[ComponentModel] | None = None + workbench: List[ComponentModel] | None = None handoffs: List[HandoffBase | str] | None = None model_context: ComponentModel | None = None memory: List[ComponentModel] | None = None @@ -73,12 +82,13 @@ class AssistantAgentConfig(BaseModel): model_client_stream: bool = False reflect_on_tool_use: bool tool_call_summary_format: str + max_tool_iterations: int = Field(default=1, ge=1) metadata: Dict[str, str] | None = None + structured_message_factory: ComponentModel | None = None class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]): """An agent that provides assistance with tool use. - The :meth:`on_messages` returns a :class:`~autogen_agentchat.base.Response` in which :attr:`~autogen_agentchat.base.Response.chat_message` is the final response message. @@ -87,10 +97,20 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]): the inner messages as they are created, and the :class:`~autogen_agentchat.base.Response` object as the last item before closing the generator. + The :meth:`BaseChatAgent.run` method returns a :class:`~autogen_agentchat.base.TaskResult` + containing the messages produced by the agent. In the list of messages, + :attr:`~autogen_agentchat.base.TaskResult.messages`, + the last message is the final response message. + + The :meth:`BaseChatAgent.run_stream` method creates an async generator that produces + the inner messages as they are created, and the :class:`~autogen_agentchat.base.TaskResult` + object as the last item before closing the generator. + .. attention:: The caller must only pass the new messages to the agent on each call - to the :meth:`on_messages` or :meth:`on_messages_stream` method. + to the :meth:`on_messages`, :meth:`on_messages_stream`, :meth:`BaseChatAgent.run`, + or :meth:`BaseChatAgent.run_stream` methods. The agent maintains its state between calls to these methods. Do not pass the entire conversation history to the agent on each call. @@ -103,21 +123,44 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]): .. image:: ../../images/assistant-agent.svg - Tool call behavior: + **Structured output:** + + If the `output_content_type` is set, the agent will respond with a :class:`~autogen_agentchat.messages.StructuredMessage` + instead of a :class:`~autogen_agentchat.messages.TextMessage` in the final response by default. - * If the model returns no tool call, then the response is immediately returned as a :class:`~autogen_agentchat.messages.TextMessage` in :attr:`~autogen_agentchat.base.Response.chat_message`. + .. note:: + + Currently, setting `output_content_type` prevents the agent from being + able to call `load_component` and `dum_component` methods for serializable + configuration. This will be fixed soon in the future. + + **Tool call behavior:** + + * If the model returns no tool call, then the response is immediately returned as a :class:`~autogen_agentchat.messages.TextMessage` or a :class:`~autogen_agentchat.messages.StructuredMessage` (when using structured output) in :attr:`~autogen_agentchat.base.Response.chat_message`. This ends the tool call iteration loop regardless of the `max_tool_iterations` setting. * When the model returns tool calls, they will be executed right away: - - When `reflect_on_tool_use` is False (default), the tool call results are returned as a :class:`~autogen_agentchat.messages.ToolCallSummaryMessage` in :attr:`~autogen_agentchat.base.Response.chat_message`. `tool_call_summary_format` can be used to customize the tool call summary. - - When `reflect_on_tool_use` is True, the another model inference is made using the tool calls and results, and the text response is returned as a :class:`~autogen_agentchat.messages.TextMessage` in :attr:`~autogen_agentchat.base.Response.chat_message`. + - When `reflect_on_tool_use` is False, the tool call results are returned as a :class:`~autogen_agentchat.messages.ToolCallSummaryMessage` in :attr:`~autogen_agentchat.base.Response.chat_message`. You can customise the summary with either a static format string (`tool_call_summary_format`) **or** a callable (`tool_call_summary_formatter`); the callable is evaluated once per tool call. + - When `reflect_on_tool_use` is True, the another model inference is made using the tool calls and results, and final response is returned as a :class:`~autogen_agentchat.messages.TextMessage` or a :class:`~autogen_agentchat.messages.StructuredMessage` (when using structured output) in :attr:`~autogen_agentchat.base.Response.chat_message`. + - `reflect_on_tool_use` is set to `True` by default when `output_content_type` is set. + - `reflect_on_tool_use` is set to `False` by default when `output_content_type` is not set. * If the model returns multiple tool calls, they will be executed concurrently. To disable parallel tool calls you need to configure the model client. For example, set `parallel_tool_calls=False` for :class:`~autogen_ext.models.openai.OpenAIChatCompletionClient` and :class:`~autogen_ext.models.openai.AzureOpenAIChatCompletionClient`. + * The `max_tool_iterations` parameter controls how many sequential tool call iterations the agent can perform in a single run. When set to 1 (default), the agent executes tool calls once and returns the result. When set higher, the agent can make additional model calls to execute more tool calls if the model continues to request them, enabling multi-step tool-based workflows. The agent stops when either the model returns a text response (instead of tool calls) or the maximum number of iterations is reached. .. tip:: - By default, the tool call results are returned as response when tool calls are made. - So it is recommended to pay attention to the formatting of the tools return values, - especially if another agent is expecting them in a specific format. - Use `tool_call_summary_format` to customize the tool call summary, if needed. - Hand off behavior: + By default, the tool call results are returned as the response when tool + calls are made, so pay close attention to how the tools' return values + are formatted—especially if another agent expects a specific schema. + + * Use **`tool_call_summary_format`** for a simple static template. + * Use **`tool_call_summary_formatter`** for full programmatic control + (e.g., "hide large success payloads, show full details on error"). + + *Note*: `tool_call_summary_formatter` is **not serializable** and will + be ignored when an agent is loaded from, or exported to, YAML/JSON + configuration files. + + + **Hand off behavior:** * If a handoff is triggered, a :class:`~autogen_agentchat.messages.HandoffMessage` will be returned in :attr:`~autogen_agentchat.base.Response.chat_message`. * If there are tool calls, they will also be executed right away before returning the handoff. @@ -129,16 +172,18 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]): To avoid this, disable parallel tool calls in the model client configuration. - Limit context size sent to the model: + **Limit context size sent to the model:** You can limit the number of messages sent to the model by setting the `model_context` parameter to a :class:`~autogen_core.model_context.BufferedChatCompletionContext`. This will limit the number of recent messages sent to the model and can be useful when the model has a limit on the number of tokens it can process. + Another option is to use a :class:`~autogen_core.model_context.TokenLimitedChatCompletionContext` + which will limit the number of tokens sent to the model. You can also create your own model context by subclassing :class:`~autogen_core.model_context.ChatCompletionContext`. - Streaming mode: + **Streaming mode:** The assistant agent can be used in streaming mode by setting `model_client_stream=True`. In this mode, the :meth:`on_messages_stream` and :meth:`BaseChatAgent.run_stream` methods will also yield @@ -146,11 +191,12 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]): messages as the model client produces chunks of response. The chunk messages will not be included in the final response's inner messages. - Args: name (str): The name of the agent. model_client (ChatCompletionClient): The model client to use for inference. tools (List[BaseTool[Any, Any] | Callable[..., Any] | Callable[..., Awaitable[Any]]] | None, optional): The tools to register with the agent. + workbench (Workbench | Sequence[Workbench] | None, optional): The workbench or list of workbenches to use for the agent. + Tools cannot be used when workbench is set and vice versa. handoffs (List[HandoffBase | str] | None, optional): The handoff configurations for the agent, allowing it to transfer to other agents by responding with a :class:`HandoffMessage`. The transfer is only executed when the team is in :class:`~autogen_agentchat.teams.Swarm`. @@ -162,13 +208,35 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]): :meth:`on_messages_stream` and :meth:`BaseChatAgent.run_stream` methods will also yield :class:`~autogen_agentchat.messages.ModelClientStreamingChunkEvent` messages as the model client produces chunks of response. Defaults to `False`. reflect_on_tool_use (bool, optional): If `True`, the agent will make another model inference using the tool call and result - to generate a response. If `False`, the tool call result will be returned as the response. Defaults to `False`. - tool_call_summary_format (str, optional): The format string used to create a tool call summary for every tool call result. - Defaults to "{result}". - When `reflect_on_tool_use` is `False`, a concatenation of all the tool call summaries, separated by a new line character ('\\n') - will be returned as the response. - Available variables: `{tool_name}`, `{arguments}`, `{result}`. - For example, `"{tool_name}: {result}"` will create a summary like `"tool_name: result"`. + to generate a response. If `False`, the tool call result will be returned as the response. By default, if `output_content_type` is set, this will be `True`; + if `output_content_type` is not set, this will be `False`. + output_content_type (type[BaseModel] | None, optional): The output content type for :class:`~autogen_agentchat.messages.StructuredMessage` response as a Pydantic model. + This will be used with the model client to generate structured output. + If this is set, the agent will respond with a :class:`~autogen_agentchat.messages.StructuredMessage` instead of a :class:`~autogen_agentchat.messages.TextMessage` + in the final response, unless `reflect_on_tool_use` is `False` and a tool call is made. + output_content_type_format (str | None, optional): (Experimental) The format string used for the content of a :class:`~autogen_agentchat.messages.StructuredMessage` response. + max_tool_iterations (int, optional): The maximum number of tool iterations to perform until the model stops making tool calls. Defaults to `1`, which means the agent will + only execute the tool calls made by the model once, and return the result as a :class:`~autogen_agentchat.messages.ToolCallSummaryMessage`, + or a :class:`~autogen_agentchat.messages.TextMessage` or a :class:`~autogen_agentchat.messages.StructuredMessage` (when using structured output) + in :attr:`~autogen_agentchat.base.Response.chat_message` as the final response. + As soon as the model stops making tool calls, the agent will stop executing tool calls and return the result as the final response. + The value must be greater than or equal to 1. + tool_call_summary_format (str, optional): Static format string applied to each tool call result when composing the :class:`~autogen_agentchat.messages.ToolCallSummaryMessage`. + Defaults to ``"{result}"``. Ignored if `tool_call_summary_formatter` is provided. When `reflect_on_tool_use` is ``False``, the summaries for all tool + calls are concatenated with a newline ('\\n') and returned as the response. Placeholders available in the template: + `{tool_name}`, `{arguments}`, `{result}`, `{is_error}`. + tool_call_summary_formatter (Callable[[FunctionCall, FunctionExecutionResult], str] | None, optional): + Callable that receives the ``FunctionCall`` and its ``FunctionExecutionResult`` and returns the summary string. + Overrides `tool_call_summary_format` when supplied and allows conditional logic — for example, emitting static string like + ``"Tool FooBar executed successfully."`` on success and a full payload (including all passed arguments etc.) only on failure. + + **Limitation**: The callable is *not serializable*; values provided via YAML/JSON configs are ignored. + + .. note:: + + `tool_call_summary_formatter` is intended for in-code use only. It cannot currently be saved or restored via + configuration files. + memory (Sequence[Memory] | None, optional): The memory store to use for the agent. Defaults to `None`. metadata (Dict[str, str] | None, optional): Optional metadata for tracking. @@ -188,10 +256,8 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]): .. code-block:: python import asyncio - from autogen_core import CancellationToken from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_agentchat.agents import AssistantAgent - from autogen_agentchat.messages import TextMessage async def main() -> None: @@ -201,10 +267,8 @@ async def main() -> None: ) agent = AssistantAgent(name="assistant", model_client=model_client) - response = await agent.on_messages( - [TextMessage(content="What is the capital of France?", source="user")], CancellationToken() - ) - print(response) + result = await agent.run(task="Name two cities in North America.") + print(result) asyncio.run(main()) @@ -219,8 +283,6 @@ async def main() -> None: import asyncio from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_agentchat.agents import AssistantAgent - from autogen_agentchat.messages import TextMessage - from autogen_core import CancellationToken async def main() -> None: @@ -234,9 +296,7 @@ async def main() -> None: model_client_stream=True, ) - stream = agent.on_messages_stream( - [TextMessage(content="Name two cities in North America.", source="user")], CancellationToken() - ) + stream = agent.run_stream(task="Name two cities in North America.") async for message in stream: print(message) @@ -245,27 +305,23 @@ async def main() -> None: .. code-block:: text - source='assistant' models_usage=None content='Two' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' cities' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' in' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' North' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' America' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' are' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' New' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' York' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' City' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' in' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' the' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' United' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' States' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' and' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' Toronto' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' in' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' Canada' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content='.' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content=' TERMIN' type='ModelClientStreamingChunkEvent' - source='assistant' models_usage=None content='ATE' type='ModelClientStreamingChunkEvent' - Response(chat_message=TextMessage(source='assistant', models_usage=RequestUsage(prompt_tokens=0, completion_tokens=0), content='Two cities in North America are New York City in the United States and Toronto in Canada. TERMINATE', type='TextMessage'), inner_messages=[]) + source='user' models_usage=None metadata={} content='Name two cities in North America.' type='TextMessage' + source='assistant' models_usage=None metadata={} content='Two' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=None metadata={} content=' cities' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=None metadata={} content=' in' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=None metadata={} content=' North' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=None metadata={} content=' America' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=None metadata={} content=' are' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=None metadata={} content=' New' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=None metadata={} content=' York' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=None metadata={} content=' City' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=None metadata={} content=' and' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=None metadata={} content=' Toronto' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=None metadata={} content='.' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=None metadata={} content=' TERMIN' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=None metadata={} content='ATE' type='ModelClientStreamingChunkEvent' + source='assistant' models_usage=RequestUsage(prompt_tokens=0, completion_tokens=0) metadata={} content='Two cities in North America are New York City and Toronto. TERMINATE' type='TextMessage' + messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Name two cities in North America.', type='TextMessage'), TextMessage(source='assistant', models_usage=RequestUsage(prompt_tokens=0, completion_tokens=0), metadata={}, content='Two cities in North America are New York City and Toronto. TERMINATE', type='TextMessage')] stop_reason=None **Example 3: agent with tools** @@ -285,9 +341,7 @@ async def main() -> None: import asyncio from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_agentchat.agents import AssistantAgent - from autogen_agentchat.messages import TextMessage from autogen_agentchat.ui import Console - from autogen_core import CancellationToken async def get_current_time() -> str: @@ -300,17 +354,102 @@ async def main() -> None: # api_key = "your_openai_api_key" ) agent = AssistantAgent(name="assistant", model_client=model_client, tools=[get_current_time]) + await Console(agent.run_stream(task="What is the current time?")) - await Console( - agent.on_messages_stream( - [TextMessage(content="What is the current time?", source="user")], CancellationToken() - ) + + asyncio.run(main()) + + **Example 4: agent with max_tool_iterations** + + The following example demonstrates how to use the `max_tool_iterations` parameter + to control how many times the agent can execute tool calls in a single run. + This is useful when you want the agent to perform multiple sequential tool + operations to reach a goal. + + .. code-block:: python + + import asyncio + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.ui import Console + + + # Global counter state + counter = 0 + + + def increment_counter() -> str: + \"\"\"Increment the counter by 1 and return the current value.\"\"\" + global counter + counter += 1 + return f"Counter incremented to: {counter}" + + + def get_counter() -> str: + \"\"\"Get the current counter value.\"\"\" + global counter + return f"Current counter value: {counter}" + + + async def main() -> None: + model_client = OpenAIChatCompletionClient( + model="gpt-4o", + # api_key = "your_openai_api_key" + ) + + # Create agent with max_tool_iterations=5 to allow multiple tool calls + agent = AssistantAgent( + name="assistant", + model_client=model_client, + tools=[increment_counter, get_counter], + max_tool_iterations=5, # Allow up to 5 tool call iterations + reflect_on_tool_use=True, # Get a final summary after tool calls + ) + + await Console(agent.run_stream(task="Increment the counter 3 times and then tell me the final value.")) + + + asyncio.run(main()) + + **Example 5: agent with Model-Context Protocol (MCP) workbench** + + The following example demonstrates how to create an assistant agent with + a model client and an :class:`~autogen_ext.tools.mcp.McpWorkbench` for + interacting with a Model-Context Protocol (MCP) server. + + .. code-block:: python + + import asyncio + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.ui import Console + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_ext.tools.mcp import StdioServerParams, McpWorkbench + + + async def main() -> None: + params = StdioServerParams( + command="uvx", + args=["mcp-server-fetch"], + read_timeout_seconds=60, ) + # You can also use `start()` and `stop()` to manage the session. + async with McpWorkbench(server_params=params) as workbench: + model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano") + assistant = AssistantAgent( + name="Assistant", + model_client=model_client, + workbench=workbench, + reflect_on_tool_use=True, + ) + await Console( + assistant.run_stream(task="Go to https://github.com/microsoft/autogen and tell me what you see.") + ) + asyncio.run(main()) - **Example 4: agent with structured output and tool** + **Example 6: agent with structured output and tool** The following example demonstrates how to create an assistant agent with a model client configured to use structured output and a tool. @@ -325,9 +464,7 @@ async def main() -> None: from typing import Literal from autogen_agentchat.agents import AssistantAgent - from autogen_agentchat.messages import TextMessage from autogen_agentchat.ui import Console - from autogen_core import CancellationToken from autogen_core.tools import FunctionTool from autogen_ext.models.openai import OpenAIChatCompletionClient from pydantic import BaseModel @@ -349,10 +486,9 @@ def sentiment_analysis(text: str) -> str: # which is required for structured output mode. tool = FunctionTool(sentiment_analysis, description="Sentiment Analysis", strict=True) - # Create an OpenAIChatCompletionClient instance that uses the structured output format. + # Create an OpenAIChatCompletionClient instance that supports structured output. model_client = OpenAIChatCompletionClient( model="gpt-4o-mini", - response_format=AgentResponse, # type: ignore ) # Create an AssistantAgent instance that uses the tool and model client. @@ -361,12 +497,12 @@ def sentiment_analysis(text: str) -> str: model_client=model_client, tools=[tool], system_message="Use the tool to analyze sentiment.", - reflect_on_tool_use=True, # Use reflection to have the agent generate a formatted response. + output_content_type=AgentResponse, ) async def main() -> None: - stream = agent.on_messages_stream([TextMessage(content="I am happy today!", source="user")], CancellationToken()) + stream = agent.run_stream(task="I am happy today!") await Console(stream) @@ -381,7 +517,7 @@ async def main() -> None: ---------- assistant ---------- {"thoughts":"The user expresses a clear positive emotion by stating they are happy today, suggesting an upbeat mood.","response":"happy"} - **Example 5: agent with bounded model context** + **Example 7: agent with bounded model context** The following example shows how to use a :class:`~autogen_core.model_context.BufferedChatCompletionContext` @@ -394,8 +530,6 @@ async def main() -> None: import asyncio from autogen_agentchat.agents import AssistantAgent - from autogen_agentchat.messages import TextMessage - from autogen_core import CancellationToken from autogen_core.model_context import BufferedChatCompletionContext from autogen_ext.models.openai import OpenAIChatCompletionClient @@ -418,20 +552,14 @@ async def main() -> None: system_message="You are a helpful assistant.", ) - response = await agent.on_messages( - [TextMessage(content="Name two cities in North America.", source="user")], CancellationToken() - ) - print(response.chat_message.content) # type: ignore + result = await agent.run(task="Name two cities in North America.") + print(result.messages[-1].content) # type: ignore - response = await agent.on_messages( - [TextMessage(content="My favorite color is blue.", source="user")], CancellationToken() - ) - print(response.chat_message.content) # type: ignore + result = await agent.run(task="My favorite color is blue.") + print(result.messages[-1].content) # type: ignore - response = await agent.on_messages( - [TextMessage(content="Did I ask you any question?", source="user")], CancellationToken() - ) - print(response.chat_message.content) # type: ignore + result = await agent.run(task="Did I ask you any question?") + print(result.messages[-1].content) # type: ignore asyncio.run(main()) @@ -442,7 +570,7 @@ async def main() -> None: That's great! Blue is often associated with calmness and serenity. Do you have a specific shade of blue that you like, or any particular reason why it's your favorite? No, you didn't ask a question. I apologize for any misunderstanding. If you have something specific you'd like to discuss or ask, feel free to let me know! - **Example 6: agent with memory** + **Example 8: agent with memory** The following example shows how to use a list-based memory with the assistant agent. The memory is preloaded with some initial content. @@ -454,8 +582,6 @@ async def main() -> None: import asyncio from autogen_agentchat.agents import AssistantAgent - from autogen_agentchat.messages import TextMessage - from autogen_core import CancellationToken from autogen_core.memory import ListMemory, MemoryContent from autogen_ext.models.openai import OpenAIChatCompletionClient @@ -480,10 +606,8 @@ async def main() -> None: system_message="You are a helpful assistant.", ) - response = await agent.on_messages( - [TextMessage(content="One idea for a dinner.", source="user")], CancellationToken() - ) - print(response.chat_message.content) # type: ignore + result = await agent.run(task="What is a good dinner idea?") + print(result.messages[-1].content) # type: ignore asyncio.run(main()) @@ -496,23 +620,21 @@ async def main() -> None: - Start with a pizza crust (store-bought or homemade). - Spread a layer of marinara or tomato sauce evenly over the crust. - Top with your favorite vegetables like bell peppers, mushrooms, onions, olives, and spinach. - - Add some protein if you’d like, such as grilled chicken or pepperoni (ensure it's cheese-free). + - Add some protein if you'd like, such as grilled chicken or pepperoni (ensure it's cheese-free). - Sprinkle with herbs like oregano and basil, and maybe a drizzle of olive oil. - Bake according to the crust instructions until the edges are golden and the veggies are cooked. Serve it with a side salad or some garlic bread to complete the meal! Enjoy your dinner! - **Example 7: agent with `o1-mini`** + **Example 9: agent with `o1-mini`** The following example shows how to use `o1-mini` model with the assistant agent. .. code-block:: python import asyncio - from autogen_core import CancellationToken from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_agentchat.agents import AssistantAgent - from autogen_agentchat.messages import TextMessage async def main() -> None: @@ -523,10 +645,8 @@ async def main() -> None: # The system message is not supported by the o1 series model. agent = AssistantAgent(name="assistant", model_client=model_client, system_message=None) - response = await agent.on_messages( - [TextMessage(content="What is the capital of France?", source="user")], CancellationToken() - ) - print(response) + result = await agent.run(task="What is the capital of France?") + print(result.messages[-1].content) # type: ignore asyncio.run(main()) @@ -538,7 +658,7 @@ async def main() -> None: See `o1 beta limitations `_ for more details. - **Example 8: agent using reasoning model with custom model context.** + **Example 10: agent using reasoning model with custom model context.** The following example shows how to use a reasoning model (DeepSeek R1) with the assistant agent. The model context is used to filter out the thought field from the assistant message. @@ -594,8 +714,10 @@ async def run_reasoning_agent() -> None: asyncio.run(run_reasoning_agent()) + For detailed examples and usage, see the Examples section below. """ + component_version = 2 component_config_schema = AssistantAgentConfig component_provider_override = "autogen_agentchat.agents.AssistantAgent" @@ -605,6 +727,7 @@ def __init__( model_client: ChatCompletionClient, *, tools: List[BaseTool[Any, Any] | Callable[..., Any] | Callable[..., Awaitable[Any]]] | None = None, + workbench: Workbench | Sequence[Workbench] | None = None, handoffs: List[HandoffBase | str] | None = None, model_context: ChatCompletionContext | None = None, description: str = "An agent that provides assistance with ability to use tools.", @@ -612,25 +735,27 @@ def __init__( str | None ) = "You are a helpful AI assistant. Solve tasks using your tools. Reply with TERMINATE when the task has been completed.", model_client_stream: bool = False, - reflect_on_tool_use: bool = False, + reflect_on_tool_use: bool | None = None, + max_tool_iterations: int = 1, tool_call_summary_format: str = "{result}", + tool_call_summary_formatter: Callable[[FunctionCall, FunctionExecutionResult], str] | None = None, + output_content_type: type[BaseModel] | None = None, + output_content_type_format: str | None = None, memory: Sequence[Memory] | None = None, metadata: Dict[str, str] | None = None, ): super().__init__(name=name, description=description) self._metadata = metadata or {} - if reflect_on_tool_use and ModelFamily.is_claude(model_client.model_info["family"]): - warnings.warn( - "Claude models may not work with reflection on tool use because Claude requires that any requests including a previous tool use or tool result must include the original tools definition." - "Consider setting reflect_on_tool_use to False. " - "As an alternative, consider calling the agent in a loop until it stops producing tool calls. " - "See [Single-Agent Team](https://microsoft.github.io/autogen/stable/user-guide/agentchat-user-guide/tutorial/teams.html#single-agent-team) " - "for more details.", - UserWarning, - stacklevel=2, - ) self._model_client = model_client self._model_client_stream = model_client_stream + self._output_content_type: type[BaseModel] | None = output_content_type + self._output_content_type_format = output_content_type_format + self._structured_message_factory: StructuredMessageFactory | None = None + if output_content_type is not None: + self._structured_message_factory = StructuredMessageFactory( + input_model=output_content_type, format_string=output_content_type_format + ) + self._memory = None if memory is not None: if isinstance(memory, list): @@ -681,49 +806,111 @@ def __init__( handoff_tool_names = [tool.name for tool in self._handoff_tools] if len(handoff_tool_names) != len(set(handoff_tool_names)): raise ValueError(f"Handoff names must be unique: {handoff_tool_names}") - # Check if handoff tool names not in tool names. - if any(name in tool_names for name in handoff_tool_names): - raise ValueError( - f"Handoff names must be unique from tool names. " - f"Handoff names: {handoff_tool_names}; tool names: {tool_names}" - ) + # Create sets for faster lookup + tool_names_set = set(tool_names) + handoff_tool_names_set = set(handoff_tool_names) + + # Check if there's any overlap between handoff tool names and tool names + overlap = tool_names_set.intersection(handoff_tool_names_set) + + # Also check if any handoff target name matches a tool name + # This handles the case where a handoff is specified directly with a string that matches a tool name + for handoff in handoffs or []: + if isinstance(handoff, str) and handoff in tool_names_set: + raise ValueError("Handoff names must be unique from tool names") + elif isinstance(handoff, HandoffBase) and handoff.target in tool_names_set: + raise ValueError("Handoff names must be unique from tool names") + + if overlap: + raise ValueError("Handoff names must be unique from tool names") + + if workbench is not None: + if self._tools: + raise ValueError("Tools cannot be used with a workbench.") + if isinstance(workbench, Sequence): + self._workbench = workbench + else: + self._workbench = [workbench] + else: + self._workbench = [StaticStreamWorkbench(self._tools)] if model_context is not None: self._model_context = model_context else: self._model_context = UnboundedChatCompletionContext() - self._reflect_on_tool_use = reflect_on_tool_use + if self._output_content_type is not None and reflect_on_tool_use is None: + # If output_content_type is set, we need to reflect on tool use by default. + self._reflect_on_tool_use = True + elif reflect_on_tool_use is None: + self._reflect_on_tool_use = False + else: + self._reflect_on_tool_use = reflect_on_tool_use + + # Tool call loop + self._max_tool_iterations = max_tool_iterations + if self._max_tool_iterations < 1: + raise ValueError( + f"Maximum number of tool iterations must be greater than or equal to 1, got {max_tool_iterations}" + ) + self._tool_call_summary_format = tool_call_summary_format + self._tool_call_summary_formatter = tool_call_summary_formatter self._is_running = False @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: - message_types: List[type[ChatMessage]] = [TextMessage] - if self._handoffs: - message_types.append(HandoffMessage) - if self._tools: - message_types.append(ToolCallSummaryMessage) - return tuple(message_types) + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: + """Get the types of messages this agent can produce. + + Returns: + Sequence of message types this agent can generate + """ + types: List[type[BaseChatMessage]] = [TextMessage, ToolCallSummaryMessage, HandoffMessage] + if self._structured_message_factory is not None: + types.append(StructuredMessage) + return types @property def model_context(self) -> ChatCompletionContext: - """ - The model context in use by the agent. + """Get the model context used by this agent. + + Returns: + The chat completion context for this agent """ return self._model_context - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages( + self, + messages: Sequence[BaseChatMessage], + cancellation_token: CancellationToken, + ) -> Response: + """Process incoming messages and generate a response. + + Args: + messages: Sequence of messages to process + cancellation_token: Token for cancelling operation + + Returns: + Response containing the agent's reply + """ async for message in self.on_messages_stream(messages, cancellation_token): if isinstance(message, Response): return message raise AssertionError("The stream should have returned the final result.") async def on_messages_stream( - self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken - ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]: - """ - Process the incoming messages with the assistant agent and yield events/responses as they happen. + self, + messages: Sequence[BaseChatMessage], + cancellation_token: CancellationToken, + ) -> AsyncGenerator[Union[BaseAgentEvent, BaseChatMessage, Response], None]: + """Process messages and stream the response. + + Args: + messages: Sequence of messages to process + cancellation_token: Token for cancelling operation + + Yields: + Events, messages and final response during processing """ # Gather all relevant state here @@ -731,13 +918,16 @@ async def on_messages_stream( model_context = self._model_context memory = self._memory system_messages = self._system_messages - tools = self._tools + workbench = self._workbench handoff_tools = self._handoff_tools handoffs = self._handoffs model_client = self._model_client model_client_stream = self._model_client_stream reflect_on_tool_use = self._reflect_on_tool_use + max_tool_iterations = self._max_tool_iterations tool_call_summary_format = self._tool_call_summary_format + tool_call_summary_formatter = self._tool_call_summary_formatter + output_content_type = self._output_content_type # STEP 1: Add new user/handoff messages to the model context await self._add_messages_to_context( @@ -746,7 +936,7 @@ async def on_messages_stream( ) # STEP 2: Update model context with any relevant memory - inner_messages: List[AgentEvent | ChatMessage] = [] + inner_messages: List[BaseAgentEvent | BaseChatMessage] = [] for event_msg in await self._update_model_context_with_memory( memory=memory, model_context=model_context, @@ -755,17 +945,22 @@ async def on_messages_stream( inner_messages.append(event_msg) yield event_msg - # STEP 3: Run the first inference + # STEP 3: Generate a message ID for correlation between streaming chunks and final message + message_id = str(uuid.uuid4()) + + # STEP 4: Run the first inference model_result = None async for inference_output in self._call_llm( model_client=model_client, model_client_stream=model_client_stream, system_messages=system_messages, model_context=model_context, - tools=tools, + workbench=workbench, handoff_tools=handoff_tools, agent_name=agent_name, cancellation_token=cancellation_token, + output_content_type=output_content_type, + message_id=message_id, ): if isinstance(inference_output, CreateResult): model_result = inference_output @@ -790,7 +985,7 @@ async def on_messages_stream( ) ) - # STEP 4: Process the model output + # STEP 5: Process the model output async for output_event in self._process_model_result( model_result=model_result, inner_messages=inner_messages, @@ -798,30 +993,34 @@ async def on_messages_stream( agent_name=agent_name, system_messages=system_messages, model_context=model_context, - tools=tools, + workbench=workbench, handoff_tools=handoff_tools, handoffs=handoffs, model_client=model_client, model_client_stream=model_client_stream, reflect_on_tool_use=reflect_on_tool_use, + max_tool_iterations=max_tool_iterations, tool_call_summary_format=tool_call_summary_format, + tool_call_summary_formatter=tool_call_summary_formatter, + output_content_type=output_content_type, + message_id=message_id, + format_string=self._output_content_type_format, ): yield output_event @staticmethod async def _add_messages_to_context( model_context: ChatCompletionContext, - messages: Sequence[ChatMessage], + messages: Sequence[BaseChatMessage], ) -> None: """ - Add incoming user (and possibly handoff) messages to the model context. + Add incoming messages to the model context. """ for msg in messages: if isinstance(msg, HandoffMessage): - # Add handoff context to the model context. - for context_msg in msg.context: - await model_context.add_message(context_msg) - await model_context.add_message(UserMessage(content=msg.content, source=msg.source)) + for llm_msg in msg.context: + await model_context.add_message(llm_msg) + await model_context.add_message(msg.to_model_message()) @staticmethod async def _update_model_context_with_memory( @@ -829,8 +1028,15 @@ async def _update_model_context_with_memory( model_context: ChatCompletionContext, agent_name: str, ) -> List[MemoryQueryEvent]: - """ - If memory modules are present, update the model context and return the events produced. + """Update model context with memory content. + + Args: + memory: Optional sequence of memory stores to query + model_context: Context to update with memory content + agent_name: Name of the agent for event tracking + + Returns: + List of memory query events generated during update """ events: List[MemoryQueryEvent] = [] if memory: @@ -851,28 +1057,47 @@ async def _call_llm( model_client_stream: bool, system_messages: List[SystemMessage], model_context: ChatCompletionContext, - tools: List[BaseTool[Any, Any]], + workbench: Sequence[Workbench], handoff_tools: List[BaseTool[Any, Any]], agent_name: str, cancellation_token: CancellationToken, + output_content_type: type[BaseModel] | None, + message_id: str, ) -> AsyncGenerator[Union[CreateResult, ModelClientStreamingChunkEvent], None]: - """ - Perform a model inference and yield either streaming chunk events or the final CreateResult. + """Call the language model with given context and configuration. + + Args: + model_client: Client for model inference + model_client_stream: Whether to stream responses + system_messages: System messages to include + model_context: Context containing message history + workbench: Available workbenches + handoff_tools: Tools for handling handoffs + agent_name: Name of the agent + cancellation_token: Token for cancelling operation + output_content_type: Optional type for structured output + + Returns: + Generator yielding model results or streaming chunks """ all_messages = await model_context.get_messages() llm_messages = cls._get_compatible_context(model_client=model_client, messages=system_messages + all_messages) - all_tools = tools + handoff_tools + tools = [tool for wb in workbench for tool in await wb.list_tools()] + handoff_tools if model_client_stream: model_result: Optional[CreateResult] = None + async for chunk in model_client.create_stream( - llm_messages, tools=all_tools, cancellation_token=cancellation_token + llm_messages, + tools=tools, + json_output=output_content_type, + cancellation_token=cancellation_token, ): if isinstance(chunk, CreateResult): model_result = chunk elif isinstance(chunk, str): - yield ModelClientStreamingChunkEvent(content=chunk, source=agent_name) + yield ModelClientStreamingChunkEvent(content=chunk, source=agent_name, full_message_id=message_id) else: raise RuntimeError(f"Invalid chunk type: {type(chunk)}") if model_result is None: @@ -880,7 +1105,10 @@ async def _call_llm( yield model_result else: model_result = await model_client.create( - llm_messages, tools=all_tools, cancellation_token=cancellation_token + llm_messages, + tools=tools, + cancellation_token=cancellation_token, + json_output=output_content_type, ) yield model_result @@ -888,97 +1116,197 @@ async def _call_llm( async def _process_model_result( cls, model_result: CreateResult, - inner_messages: List[AgentEvent | ChatMessage], + inner_messages: List[BaseAgentEvent | BaseChatMessage], cancellation_token: CancellationToken, agent_name: str, system_messages: List[SystemMessage], model_context: ChatCompletionContext, - tools: List[BaseTool[Any, Any]], + workbench: Sequence[Workbench], handoff_tools: List[BaseTool[Any, Any]], handoffs: Dict[str, HandoffBase], model_client: ChatCompletionClient, model_client_stream: bool, reflect_on_tool_use: bool, tool_call_summary_format: str, - ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]: + tool_call_summary_formatter: Callable[[FunctionCall, FunctionExecutionResult], str] | None, + max_tool_iterations: int, + output_content_type: type[BaseModel] | None, + message_id: str, + format_string: str | None = None, + ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]: """ Handle final or partial responses from model_result, including tool calls, handoffs, - and reflection if needed. + and reflection if needed. Supports tool call loops when enabled. """ - # If direct text response (string) - if isinstance(model_result.content, str): - yield Response( - chat_message=TextMessage( - content=model_result.content, - source=agent_name, - models_usage=model_result.usage, - ), - inner_messages=inner_messages, - ) - return + # Tool call loop implementation with streaming support + current_model_result = model_result + # This variable is needed for the final summary/reflection step + executed_calls_and_results: List[Tuple[FunctionCall, FunctionExecutionResult]] = [] + + for loop_iteration in range(max_tool_iterations): + # If direct text response (string), we're done + if isinstance(current_model_result.content, str): + # Use the passed message ID for the final message + if output_content_type: + content = output_content_type.model_validate_json(current_model_result.content) + yield Response( + chat_message=StructuredMessage[output_content_type]( # type: ignore[valid-type] + content=content, + source=agent_name, + models_usage=current_model_result.usage, + format_string=format_string, + id=message_id, + ), + inner_messages=inner_messages, + ) + else: + yield Response( + chat_message=TextMessage( + content=current_model_result.content, + source=agent_name, + models_usage=current_model_result.usage, + id=message_id, + ), + inner_messages=inner_messages, + ) + return - # Otherwise, we have function calls - assert isinstance(model_result.content, list) and all( - isinstance(item, FunctionCall) for item in model_result.content - ) + # Otherwise, we have function calls + assert isinstance(current_model_result.content, list) and all( + isinstance(item, FunctionCall) for item in current_model_result.content + ) - # STEP 4A: Yield ToolCallRequestEvent - tool_call_msg = ToolCallRequestEvent( - content=model_result.content, - source=agent_name, - models_usage=model_result.usage, - ) - event_logger.debug(tool_call_msg) - inner_messages.append(tool_call_msg) - yield tool_call_msg - - # STEP 4B: Execute tool calls - executed_calls_and_results = await asyncio.gather( - *[ - cls._execute_tool_call( - tool_call=call, - tools=tools, - handoff_tools=handoff_tools, - agent_name=agent_name, - cancellation_token=cancellation_token, + # STEP 4A: Yield ToolCallRequestEvent + tool_call_msg = ToolCallRequestEvent( + content=current_model_result.content, + source=agent_name, + models_usage=current_model_result.usage, + ) + event_logger.debug(tool_call_msg) + inner_messages.append(tool_call_msg) + yield tool_call_msg + + # STEP 4B: Execute tool calls with streaming support + # Use a queue to handle streaming results from tool calls. + stream = asyncio.Queue[BaseAgentEvent | BaseChatMessage | None]() + + async def _execute_tool_calls( + function_calls: List[FunctionCall], + stream_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | None], + ) -> List[Tuple[FunctionCall, FunctionExecutionResult]]: + results = await asyncio.gather( + *[ + cls._execute_tool_call( + tool_call=call, + workbench=workbench, + handoff_tools=handoff_tools, + agent_name=agent_name, + cancellation_token=cancellation_token, + stream=stream_queue, + ) + for call in function_calls + ] ) - for call in model_result.content - ] - ) - exec_results = [result for _, result in executed_calls_and_results] + # Signal the end of streaming by putting None in the queue. + stream_queue.put_nowait(None) + return results + + task = asyncio.create_task(_execute_tool_calls(current_model_result.content, stream)) + + while True: + event = await stream.get() + if event is None: + # End of streaming, break the loop. + break + if isinstance(event, BaseAgentEvent) or isinstance(event, BaseChatMessage): + yield event + inner_messages.append(event) + else: + raise RuntimeError(f"Unexpected event type: {type(event)}") - # Yield ToolCallExecutionEvent - tool_call_result_msg = ToolCallExecutionEvent( - content=exec_results, - source=agent_name, - ) - event_logger.debug(tool_call_result_msg) - await model_context.add_message(FunctionExecutionResultMessage(content=exec_results)) - inner_messages.append(tool_call_result_msg) - yield tool_call_result_msg + # Wait for all tool calls to complete. + executed_calls_and_results = await task + exec_results = [result for _, result in executed_calls_and_results] - # STEP 4C: Check for handoff - handoff_output = cls._check_and_handle_handoff( - model_result=model_result, - executed_calls_and_results=executed_calls_and_results, - inner_messages=inner_messages, - handoffs=handoffs, - agent_name=agent_name, - ) - if handoff_output: - yield handoff_output - return + # Yield ToolCallExecutionEvent + tool_call_result_msg = ToolCallExecutionEvent( + content=exec_results, + source=agent_name, + ) + event_logger.debug(tool_call_result_msg) + await model_context.add_message(FunctionExecutionResultMessage(content=exec_results)) + inner_messages.append(tool_call_result_msg) + yield tool_call_result_msg + + # STEP 4C: Check for handoff + handoff_output = cls._check_and_handle_handoff( + model_result=current_model_result, + executed_calls_and_results=executed_calls_and_results, + inner_messages=inner_messages, + handoffs=handoffs, + agent_name=agent_name, + ) + if handoff_output: + yield handoff_output + return + + # STEP 4D: Check if we should continue the loop. + # If we are on the last iteration, break to the summary/reflection step. + if loop_iteration == max_tool_iterations - 1: + break + + # Continue the loop: make another model call using _call_llm + next_model_result: Optional[CreateResult] = None + async for llm_output in cls._call_llm( + model_client=model_client, + model_client_stream=model_client_stream, + system_messages=system_messages, + model_context=model_context, + workbench=workbench, + handoff_tools=handoff_tools, + agent_name=agent_name, + cancellation_token=cancellation_token, + output_content_type=output_content_type, + message_id=message_id, # Use same message ID for consistency + ): + if isinstance(llm_output, CreateResult): + next_model_result = llm_output + else: + # Streaming chunk event + yield llm_output + + assert next_model_result is not None, "No model result was produced in tool call loop." + current_model_result = next_model_result + + # Yield thought event if present + if current_model_result.thought: + thought_event = ThoughtEvent(content=current_model_result.thought, source=agent_name) + yield thought_event + inner_messages.append(thought_event) + + # Add the assistant message to the model context (including thought if present) + await model_context.add_message( + AssistantMessage( + content=current_model_result.content, + source=agent_name, + thought=getattr(current_model_result, "thought", None), + ) + ) - # STEP 4D: Reflect or summarize tool results + # After the loop, reflect or summarize tool results if reflect_on_tool_use: async for reflection_response in cls._reflect_on_tool_use_flow( system_messages=system_messages, model_client=model_client, model_client_stream=model_client_stream, model_context=model_context, + workbench=workbench, + handoff_tools=handoff_tools, agent_name=agent_name, inner_messages=inner_messages, + output_content_type=output_content_type, + cancellation_token=cancellation_token, ): yield reflection_response else: @@ -987,20 +1315,30 @@ async def _process_model_result( inner_messages=inner_messages, handoffs=handoffs, tool_call_summary_format=tool_call_summary_format, + tool_call_summary_formatter=tool_call_summary_formatter, agent_name=agent_name, ) + return @staticmethod def _check_and_handle_handoff( model_result: CreateResult, executed_calls_and_results: List[Tuple[FunctionCall, FunctionExecutionResult]], - inner_messages: List[AgentEvent | ChatMessage], + inner_messages: List[BaseAgentEvent | BaseChatMessage], handoffs: Dict[str, HandoffBase], agent_name: str, ) -> Optional[Response]: - """ - Detect handoff calls, generate the HandoffMessage if needed, and return a Response. - If multiple handoffs exist, only the first is used. + """Check for and handle any handoff requests in the model result. + + Args: + model_result: Result from model inference + executed_calls_and_results: List of executed tool calls and their results + inner_messages: List of messages generated during processing + handoffs: Dictionary of available handoff configurations + agent_name: Name of the agent + + Returns: + Optional response containing handoff message if handoff detected """ handoff_reqs = [ call for call in model_result.content if isinstance(call, FunctionCall) and call.name in handoffs @@ -1042,6 +1380,14 @@ def _check_and_handle_handoff( ) ) handoff_context.append(FunctionExecutionResultMessage(content=tool_call_results)) + elif model_result.thought: + # If no tool calls, but a thought exists, include it in the context + handoff_context.append( + AssistantMessage( + content=model_result.thought, + source=agent_name, + ) + ) # Return response for the first handoff return Response( @@ -1062,8 +1408,12 @@ async def _reflect_on_tool_use_flow( model_client: ChatCompletionClient, model_client_stream: bool, model_context: ChatCompletionContext, + workbench: Sequence[Workbench], + handoff_tools: List[BaseTool[Any, Any]], agent_name: str, - inner_messages: List[AgentEvent | ChatMessage], + inner_messages: List[BaseAgentEvent | BaseChatMessage], + output_content_type: type[BaseModel] | None, + cancellation_token: CancellationToken, ) -> AsyncGenerator[Response | ModelClientStreamingChunkEvent | ThoughtEvent, None]: """ If reflect_on_tool_use=True, we do another inference based on tool results @@ -1074,16 +1424,31 @@ async def _reflect_on_tool_use_flow( reflection_result: Optional[CreateResult] = None + # Generate a message ID for correlation between chunks and final message in reflection flow + reflection_message_id = str(uuid.uuid4()) + if model_client_stream: - async for chunk in model_client.create_stream(llm_messages): + async for chunk in model_client.create_stream( + llm_messages, + json_output=output_content_type, + cancellation_token=cancellation_token, + tool_choice="none", # Do not use tools in reflection flow. + ): if isinstance(chunk, CreateResult): reflection_result = chunk elif isinstance(chunk, str): - yield ModelClientStreamingChunkEvent(content=chunk, source=agent_name) + yield ModelClientStreamingChunkEvent( + content=chunk, source=agent_name, full_message_id=reflection_message_id + ) else: raise RuntimeError(f"Invalid chunk type: {type(chunk)}") else: - reflection_result = await model_client.create(llm_messages) + reflection_result = await model_client.create( + llm_messages, + json_output=output_content_type, + cancellation_token=cancellation_token, + tool_choice="none", # Do not use tools in reflection flow. + ) if not reflection_result or not isinstance(reflection_result.content, str): raise RuntimeError("Reflect on tool use produced no valid text response.") @@ -1103,21 +1468,35 @@ async def _reflect_on_tool_use_flow( ) ) - yield Response( - chat_message=TextMessage( - content=reflection_result.content, - source=agent_name, - models_usage=reflection_result.usage, - ), - inner_messages=inner_messages, - ) + if output_content_type: + content = output_content_type.model_validate_json(reflection_result.content) + yield Response( + chat_message=StructuredMessage[output_content_type]( # type: ignore[valid-type] + content=content, + source=agent_name, + models_usage=reflection_result.usage, + id=reflection_message_id, + ), + inner_messages=inner_messages, + ) + else: + yield Response( + chat_message=TextMessage( + content=reflection_result.content, + source=agent_name, + models_usage=reflection_result.usage, + id=reflection_message_id, + ), + inner_messages=inner_messages, + ) @staticmethod def _summarize_tool_use( executed_calls_and_results: List[Tuple[FunctionCall, FunctionExecutionResult]], - inner_messages: List[AgentEvent | ChatMessage], + inner_messages: List[BaseAgentEvent | BaseChatMessage], handoffs: Dict[str, HandoffBase], tool_call_summary_format: str, + tool_call_summary_formatter: Callable[[FunctionCall, FunctionExecutionResult], str] | None, agent_name: str, ) -> Response: """ @@ -1125,20 +1504,26 @@ def _summarize_tool_use( """ # Filter out calls which were actually handoffs normal_tool_calls = [(call, result) for call, result in executed_calls_and_results if call.name not in handoffs] - tool_call_summaries: List[str] = [] - for tool_call, tool_call_result in normal_tool_calls: - tool_call_summaries.append( - tool_call_summary_format.format( - tool_name=tool_call.name, - arguments=tool_call.arguments, - result=tool_call_result.content, - ) + + def default_tool_call_summary_formatter(call: FunctionCall, result: FunctionExecutionResult) -> str: + return tool_call_summary_format.format( + tool_name=call.name, + arguments=call.arguments, + result=result.content, + is_error=result.is_error, ) + + summary_formatter = tool_call_summary_formatter or default_tool_call_summary_formatter + + tool_call_summaries = [summary_formatter(call, result) for call, result in normal_tool_calls] + tool_call_summary = "\n".join(tool_call_summaries) return Response( chat_message=ToolCallSummaryMessage( content=tool_call_summary, source=agent_name, + tool_calls=[call for call, _ in normal_tool_calls], + results=[result for _, result in normal_tool_calls], ), inner_messages=inner_messages, ) @@ -1146,32 +1531,17 @@ def _summarize_tool_use( @staticmethod async def _execute_tool_call( tool_call: FunctionCall, - tools: List[BaseTool[Any, Any]], + workbench: Sequence[Workbench], handoff_tools: List[BaseTool[Any, Any]], agent_name: str, cancellation_token: CancellationToken, + stream: asyncio.Queue[BaseAgentEvent | BaseChatMessage | None], ) -> Tuple[FunctionCall, FunctionExecutionResult]: """Execute a single tool call and return the result.""" + # Load the arguments from the tool call. try: - all_tools = tools + handoff_tools - if not all_tools: - raise ValueError("No tools are available.") - tool = next((t for t in all_tools if t.name == tool_call.name), None) - if tool is None: - raise ValueError(f"The tool '{tool_call.name}' is not available.") - arguments: Dict[str, Any] = json.loads(tool_call.arguments) if tool_call.arguments else {} - result = await tool.run_json(arguments, cancellation_token) - result_as_str = tool.return_value_as_string(result) - return ( - tool_call, - FunctionExecutionResult( - content=result_as_str, - call_id=tool_call.id, - is_error=False, - name=tool_call.name, - ), - ) - except Exception as e: + arguments = json.loads(tool_call.arguments) + except json.JSONDecodeError as e: return ( tool_call, FunctionExecutionResult( @@ -1182,6 +1552,73 @@ async def _execute_tool_call( ), ) + # Check if the tool call is a handoff. + # TODO: consider creating a combined workbench to handle both handoff and normal tools. + for handoff_tool in handoff_tools: + if tool_call.name == handoff_tool.name: + # Run handoff tool call. + result = await handoff_tool.run_json(arguments, cancellation_token, call_id=tool_call.id) + result_as_str = handoff_tool.return_value_as_string(result) + return ( + tool_call, + FunctionExecutionResult( + content=result_as_str, + call_id=tool_call.id, + is_error=False, + name=tool_call.name, + ), + ) + + # Handle normal tool call using workbench. + for wb in workbench: + tools = await wb.list_tools() + if any(t["name"] == tool_call.name for t in tools): + if isinstance(wb, StaticStreamWorkbench): + tool_result: ToolResult | None = None + async for event in wb.call_tool_stream( + name=tool_call.name, + arguments=arguments, + cancellation_token=cancellation_token, + call_id=tool_call.id, + ): + if isinstance(event, ToolResult): + tool_result = event + elif isinstance(event, BaseAgentEvent) or isinstance(event, BaseChatMessage): + await stream.put(event) + else: + warnings.warn( + f"Unexpected event type: {type(event)} in tool call streaming.", + UserWarning, + stacklevel=2, + ) + assert isinstance(tool_result, ToolResult), "Tool result should not be None in streaming mode." + else: + tool_result = await wb.call_tool( + name=tool_call.name, + arguments=arguments, + cancellation_token=cancellation_token, + call_id=tool_call.id, + ) + return ( + tool_call, + FunctionExecutionResult( + content=tool_result.to_text(), + call_id=tool_call.id, + is_error=tool_result.is_error, + name=tool_call.name, + ), + ) + + return ( + tool_call, + FunctionExecutionResult( + content=f"Error: tool '{tool_call.name}' not found in any workbench", + call_id=tool_call.id, + is_error=True, + name=tool_call.name, + ), + ) + async def on_reset(self, cancellation_token: CancellationToken) -> None: """Reset the assistant agent to its initialization state.""" await self._model_context.clear() @@ -1211,7 +1648,8 @@ def _to_config(self) -> AssistantAgentConfig: return AssistantAgentConfig( name=self.name, model_client=self._model_client.dump_component(), - tools=[tool.dump_component() for tool in self._tools], + tools=None, # versionchanged:: v0.5.5 Now tools are not serialized, Cause they are part of the workbench. + workbench=[wb.dump_component() for wb in self._workbench] if self._workbench else None, handoffs=list(self._handoffs.values()) if self._handoffs else None, model_context=self._model_context.dump_component(), memory=[memory.dump_component() for memory in self._memory] if self._memory else None, @@ -1221,24 +1659,41 @@ def _to_config(self) -> AssistantAgentConfig: else None, model_client_stream=self._model_client_stream, reflect_on_tool_use=self._reflect_on_tool_use, + max_tool_iterations=self._max_tool_iterations, tool_call_summary_format=self._tool_call_summary_format, + structured_message_factory=self._structured_message_factory.dump_component() + if self._structured_message_factory + else None, metadata=self._metadata, ) @classmethod def _from_config(cls, config: AssistantAgentConfig) -> Self: """Create an assistant agent from a declarative config.""" + if config.structured_message_factory: + structured_message_factory = StructuredMessageFactory.load_component(config.structured_message_factory) + format_string = structured_message_factory.format_string + output_content_type = structured_message_factory.ContentModel + + else: + format_string = None + output_content_type = None + return cls( name=config.name, model_client=ChatCompletionClient.load_component(config.model_client), - tools=[BaseTool.load_component(tool) for tool in config.tools] if config.tools else None, + workbench=[Workbench.load_component(wb) for wb in config.workbench] if config.workbench else None, handoffs=config.handoffs, - model_context=None, + model_context=ChatCompletionContext.load_component(config.model_context) if config.model_context else None, + tools=[BaseTool.load_component(tool) for tool in config.tools] if config.tools else None, memory=[Memory.load_component(memory) for memory in config.memory] if config.memory else None, description=config.description, system_message=config.system_message, model_client_stream=config.model_client_stream, reflect_on_tool_use=config.reflect_on_tool_use, + max_tool_iterations=config.max_tool_iterations, tool_call_summary_format=config.tool_call_summary_format, + output_content_type=output_content_type, + output_content_type_format=format_string, metadata=config.metadata, ) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_base_chat_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_base_chat_agent.py index 94b89235df89..ea4a74a28e62 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_base_chat_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_base_chat_agent.py @@ -1,14 +1,13 @@ from abc import ABC, abstractmethod from typing import Any, AsyncGenerator, List, Mapping, Sequence -from autogen_core import CancellationToken, ComponentBase +from autogen_core import CancellationToken, ComponentBase, trace_create_agent_span, trace_invoke_agent_span from pydantic import BaseModel from ..base import ChatAgent, Response, TaskResult from ..messages import ( - AgentEvent, + BaseAgentEvent, BaseChatMessage, - ChatMessage, ModelClientStreamingChunkEvent, TextMessage, ) @@ -40,10 +39,15 @@ class BaseChatAgent(ChatAgent, ABC, ComponentBase[BaseModel]): component_type = "agent" def __init__(self, name: str, description: str) -> None: - self._name = name - if self._name.isidentifier() is False: - raise ValueError("The agent name must be a valid Python identifier.") - self._description = description + """Initialize the agent with a name and description.""" + with trace_create_agent_span( + agent_name=name, + agent_description=description, + ): + self._name = name + if self._name.isidentifier() is False: + raise ValueError("The agent name must be a valid Python identifier.") + self._description = description @property def name(self) -> str: @@ -60,13 +64,13 @@ def description(self) -> str: @property @abstractmethod - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: """The types of messages that the agent produces in the - :attr:`Response.chat_message` field. They must be :class:`ChatMessage` types.""" + :attr:`Response.chat_message` field. They must be :class:`BaseChatMessage` types.""" ... @abstractmethod - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: """Handles incoming messages and returns a response. .. note:: @@ -82,8 +86,8 @@ async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: ... async def on_messages_stream( - self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken - ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]: + self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken + ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]: """Handles incoming messages and returns a stream of messages and and the final item is the response. The base implementation in :class:`BaseChatAgent` simply calls :meth:`on_messages` and yields @@ -107,83 +111,105 @@ async def on_messages_stream( async def run( self, *, - task: str | ChatMessage | Sequence[ChatMessage] | None = None, + task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None, cancellation_token: CancellationToken | None = None, + output_task_messages: bool = True, ) -> TaskResult: """Run the agent with the given task and return the result.""" - if cancellation_token is None: - cancellation_token = CancellationToken() - input_messages: List[ChatMessage] = [] - output_messages: List[AgentEvent | ChatMessage] = [] - if task is None: - pass - elif isinstance(task, str): - text_msg = TextMessage(content=task, source="user") - input_messages.append(text_msg) - output_messages.append(text_msg) - elif isinstance(task, BaseChatMessage): - input_messages.append(task) - output_messages.append(task) - else: - if not task: - raise ValueError("Task list cannot be empty.") - # Task is a sequence of messages. - for msg in task: - if isinstance(msg, BaseChatMessage): - input_messages.append(msg) - output_messages.append(msg) - else: - raise ValueError(f"Invalid message type in sequence: {type(msg)}") - response = await self.on_messages(input_messages, cancellation_token) - if response.inner_messages is not None: - output_messages += response.inner_messages - output_messages.append(response.chat_message) - return TaskResult(messages=output_messages) + with trace_invoke_agent_span( + agent_name=self.name, + agent_description=self.description, + ): + if cancellation_token is None: + cancellation_token = CancellationToken() + input_messages: List[BaseChatMessage] = [] + output_messages: List[BaseAgentEvent | BaseChatMessage] = [] + if task is None: + pass + elif isinstance(task, str): + text_msg = TextMessage(content=task, source="user") + input_messages.append(text_msg) + if output_task_messages: + output_messages.append(text_msg) + elif isinstance(task, BaseChatMessage): + input_messages.append(task) + if output_task_messages: + output_messages.append(task) + else: + if not task: + raise ValueError("Task list cannot be empty.") + # Task is a sequence of messages. + for msg in task: + if isinstance(msg, BaseChatMessage): + input_messages.append(msg) + if output_task_messages: + output_messages.append(msg) + else: + raise ValueError(f"Invalid message type in sequence: {type(msg)}") + response = await self.on_messages(input_messages, cancellation_token) + if response.inner_messages is not None: + output_messages += response.inner_messages + output_messages.append(response.chat_message) + return TaskResult(messages=output_messages) async def run_stream( self, *, - task: str | ChatMessage | Sequence[ChatMessage] | None = None, + task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None, cancellation_token: CancellationToken | None = None, - ) -> AsyncGenerator[AgentEvent | ChatMessage | TaskResult, None]: + output_task_messages: bool = True, + ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None]: """Run the agent with the given task and return a stream of messages - and the final task result as the last item in the stream.""" - if cancellation_token is None: - cancellation_token = CancellationToken() - input_messages: List[ChatMessage] = [] - output_messages: List[AgentEvent | ChatMessage] = [] - if task is None: - pass - elif isinstance(task, str): - text_msg = TextMessage(content=task, source="user") - input_messages.append(text_msg) - output_messages.append(text_msg) - yield text_msg - elif isinstance(task, BaseChatMessage): - input_messages.append(task) - output_messages.append(task) - yield task - else: - if not task: - raise ValueError("Task list cannot be empty.") - for msg in task: - if isinstance(msg, BaseChatMessage): - input_messages.append(msg) - output_messages.append(msg) - yield msg - else: - raise ValueError(f"Invalid message type in sequence: {type(msg)}") - async for message in self.on_messages_stream(input_messages, cancellation_token): - if isinstance(message, Response): - yield message.chat_message - output_messages.append(message.chat_message) - yield TaskResult(messages=output_messages) + and the final task result as the last item in the stream. + + Args: + task: The task to run. Can be a string, a single message, or a sequence of messages. + cancellation_token: The cancellation token to kill the task immediately. + output_task_messages: Whether to include task messages in the output stream. Defaults to True for backward compatibility. + """ + with trace_invoke_agent_span( + agent_name=self.name, + agent_description=self.description, + ): + if cancellation_token is None: + cancellation_token = CancellationToken() + input_messages: List[BaseChatMessage] = [] + output_messages: List[BaseAgentEvent | BaseChatMessage] = [] + if task is None: + pass + elif isinstance(task, str): + text_msg = TextMessage(content=task, source="user") + input_messages.append(text_msg) + if output_task_messages: + output_messages.append(text_msg) + yield text_msg + elif isinstance(task, BaseChatMessage): + input_messages.append(task) + if output_task_messages: + output_messages.append(task) + yield task else: - yield message - if isinstance(message, ModelClientStreamingChunkEvent): - # Skip the model client streaming chunk events. - continue - output_messages.append(message) + if not task: + raise ValueError("Task list cannot be empty.") + for msg in task: + if isinstance(msg, BaseChatMessage): + input_messages.append(msg) + if output_task_messages: + output_messages.append(msg) + yield msg + else: + raise ValueError(f"Invalid message type in sequence: {type(msg)}") + async for message in self.on_messages_stream(input_messages, cancellation_token): + if isinstance(message, Response): + yield message.chat_message + output_messages.append(message.chat_message) + yield TaskResult(messages=output_messages) + else: + yield message + if isinstance(message, ModelClientStreamingChunkEvent): + # Skip the model client streaming chunk events. + continue + output_messages.append(message) @abstractmethod async def on_reset(self, cancellation_token: CancellationToken) -> None: diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py index 089daa2a15f0..4e3fcf0e9ec3 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py @@ -1,43 +1,146 @@ +import logging import re -from typing import List, Sequence +from inspect import iscoroutinefunction +from typing import ( + AsyncGenerator, + Awaitable, + Callable, + List, + Optional, + Sequence, + Union, + cast, +) from autogen_core import CancellationToken, Component, ComponentModel -from autogen_core.code_executor import CodeBlock, CodeExecutor +from autogen_core.code_executor import CodeBlock, CodeExecutor, CodeResult +from autogen_core.model_context import ( + ChatCompletionContext, + UnboundedChatCompletionContext, +) +from autogen_core.models import ( + AssistantMessage, + ChatCompletionClient, + CreateResult, + LLMMessage, + SystemMessage, + UserMessage, +) from pydantic import BaseModel from typing_extensions import Self +from .. import EVENT_LOGGER_NAME from ..base import Response -from ..messages import ChatMessage, TextMessage +from ..messages import ( + BaseAgentEvent, + BaseChatMessage, + CodeExecutionEvent, + CodeGenerationEvent, + HandoffMessage, + ModelClientStreamingChunkEvent, + TextMessage, + ThoughtEvent, +) +from ..utils import remove_images from ._base_chat_agent import BaseChatAgent +event_logger = logging.getLogger(EVENT_LOGGER_NAME) + class CodeExecutorAgentConfig(BaseModel): """Configuration for CodeExecutorAgent""" name: str code_executor: ComponentModel - description: str = "A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks)." + model_client: ComponentModel | None = None + description: str | None = None sources: List[str] | None = None + system_message: str | None = None + model_client_stream: bool = False + model_context: ComponentModel | None = None + supported_languages: List[str] | None = None + + +class RetryDecision(BaseModel): + reason: str + retry: bool + + +class ApprovalRequest(BaseModel): + """Request for approval of code execution.""" + + code: str + context: List[LLMMessage] + + +class ApprovalResponse(BaseModel): + """Response to approval request.""" + + approved: bool + reason: str + + +# Type aliases for approval functions +SyncApprovalFunc = Callable[[ApprovalRequest], ApprovalResponse] +AsyncApprovalFunc = Callable[[ApprovalRequest], Awaitable[ApprovalResponse]] +ApprovalFuncType = Union[SyncApprovalFunc, AsyncApprovalFunc] class CodeExecutorAgent(BaseChatAgent, Component[CodeExecutorAgentConfig]): - """An agent that extracts and executes code snippets found in received messages and returns the output. + """(Experimental) An agent that generates and executes code snippets based on user instructions. - It is typically used within a team with another agent that generates code snippets to be executed. + .. note:: + + This agent is experimental and may change in future releases. + + It is typically used within a team with another agent that generates code snippets + to be executed or alone with `model_client` provided so that it can generate code + based on user query, execute it and reflect on the code result. + + When used with `model_client`, it will generate code snippets using the model + and execute them using the provided `code_executor`. The model will also reflect on the + code execution results. The agent will yield the final reflection result from the model + as the final response. + + When used without `model_client`, it will only execute code blocks found in + :class:`~autogen_agentchat.messages.TextMessage` messages and returns the output + of the code execution. .. note:: - Consider :class:`~autogen_ext.tools.code_execution.PythonCodeExecutionTool` - as an alternative to this agent. The tool allows for executing Python code - within a single agent, rather than sending it to a separate agent for execution. - However, the model for the agent will have to generate properly escaped code - string as a parameter to the tool. + Using :class:`~autogen_agentchat.agents.AssistantAgent` with + :class:`~autogen_ext.tools.code_execution.PythonCodeExecutionTool` + is an alternative to this agent. However, the model for that agent will + have to generate properly escaped code string as a parameter to the tool. Args: - name: The name of the agent. - code_executor: The CodeExecutor responsible for executing code received in messages (:py:class:`~autogen_ext.code_executors.docker.DockerCommandLineCodeExecutor` recommended. See example below) - description (optional): The description of the agent. - sources (optional): Check only messages from the specified agents for the code to execute. + name (str): The name of the agent. + code_executor (CodeExecutor): The code executor responsible for executing code received in messages + (:py:class:`~autogen_ext.code_executors.docker.DockerCommandLineCodeExecutor` recommended. See example below) + model_client (ChatCompletionClient, optional): The model client to use for inference and generating code. + If not provided, the agent will only execute code blocks found in input messages. + Currently, the model must support structured output mode, which is required for + the automatic retry mechanism to work. + model_client_stream (bool, optional): If `True`, the model client will be used in streaming mode. + :meth:`on_messages_stream` and :meth:`BaseChatAgent.run_stream` methods will + also yield :class:`~autogen_agentchat.messages.ModelClientStreamingChunkEvent` + messages as the model client produces chunks of response. Defaults to `False`. + description (str, optional): The description of the agent. If not provided, + :class:`~autogen_agentchat.agents.CodeExecutorAgent.DEFAULT_AGENT_DESCRIPTION` will be used. + system_message (str, optional): The system message for the model. If provided, it will be prepended to the messages in the model context when making an inference. Set to `None` to disable. + Defaults to :class:`~autogen_agentchat.agents.CodeExecutorAgent.DEFAULT_SYSTEM_MESSAGE`. This is only used if `model_client` is provided. + sources (Sequence[str], optional): Check only messages from the specified agents for the code to execute. + This is useful when the agent is part of a group chat and you want to limit the code execution to messages from specific agents. + If not provided, all messages will be checked for code blocks. + This is only used if `model_client` is not provided. + max_retries_on_error (int, optional): The maximum number of retries on error. If the code execution fails, the agent will retry up to this number of times. + If the code execution fails after this number of retries, the agent will yield a reflection result. + supported_languages (List[str], optional): List of programming languages that will be parsed and executed from agent response; + others will be ignored. Defaults to DEFAULT_SUPPORTED_LANGUAGES. + approval_func (Optional[Union[Callable[[ApprovalRequest], ApprovalResponse], Callable[[ApprovalRequest], Awaitable[ApprovalResponse]]]], optional): A function that is called before each code execution to get approval. + The function takes an ApprovalRequest containing the code to be executed and the current context, and returns an ApprovalResponse. + The function can be either synchronous or asynchronous. If None (default), all code executions are automatically approved. + If set, the agent cannot be serialized using :meth:`~autogen_agentchat.agents.CodeExecutorAgent.dump_component`. .. note:: @@ -64,22 +167,44 @@ class CodeExecutorAgent(BaseChatAgent, Component[CodeExecutorAgentConfig]): In this example, we show how to set up a `CodeExecutorAgent` agent that uses the :py:class:`~autogen_ext.code_executors.docker.DockerCommandLineCodeExecutor` - to execute code snippets in a Docker container. The `work_dir` parameter indicates where all executed files are first saved locally before being executed in the Docker container. + to execute code snippets in a Docker container. The `work_dir` parameter indicates + where all executed files are first saved locally before being executed in the Docker container. .. code-block:: python import asyncio - from autogen_agentchat.agents import CodeExecutorAgent + from autogen_agentchat.agents import CodeExecutorAgent, ApprovalRequest, ApprovalResponse from autogen_agentchat.messages import TextMessage from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor from autogen_core import CancellationToken + def simple_approval_func(request: ApprovalRequest) -> ApprovalResponse: + \"\"\"Simple approval function that requests user input for code execution approval.\"\"\" + print("Code execution approval requested:") + print("=" * 50) + print(request.code) + print("=" * 50) + + while True: + user_input = input("Do you want to execute this code? (y/n): ").strip().lower() + if user_input in ['y', 'yes']: + return ApprovalResponse(approved=True, reason='Approved by user') + elif user_input in ['n', 'no']: + return ApprovalResponse(approved=False, reason='Denied by user') + else: + print("Please enter 'y' for yes or 'n' for no.") + + async def run_code_executor_agent() -> None: # Create a code executor agent that uses a Docker container to execute code. code_executor = DockerCommandLineCodeExecutor(work_dir="coding") await code_executor.start() - code_executor_agent = CodeExecutorAgent("code_executor", code_executor=code_executor) + code_executor_agent = CodeExecutorAgent( + "code_executor", + code_executor=code_executor, + approval_func=simple_approval_func + ) # Run the agent with a given code snippet. task = TextMessage( @@ -99,8 +224,204 @@ async def run_code_executor_agent() -> None: asyncio.run(run_code_executor_agent()) + In this example, we show how to set up a `CodeExecutorAgent` agent that uses the + :py:class:`~docker.types.DeviceRequest` to expose a GPU to the container for cuda-accelerated code execution. + + .. code-block:: python + + import asyncio + from autogen_agentchat.agents import CodeExecutorAgent + from autogen_agentchat.messages import TextMessage + from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor + from autogen_core import CancellationToken + from docker.types import DeviceRequest + + + async def run_code_executor_agent() -> None: + # Create a code executor agent that uses a Docker container to execute code. + code_executor = DockerCommandLineCodeExecutor( + work_dir="coding", device_requests=[DeviceRequest(count=-1, capabilities=[["gpu"]])] + ) + await code_executor.start() + code_executor_agent = CodeExecutorAgent("code_executor", code_executor=code_executor) + + # Display the GPU information + task = TextMessage( + content='''Here is some code + ```sh + nvidia-smi + ``` + ''', + source="user", + ) + response = await code_executor_agent.on_messages([task], CancellationToken()) + print(response.chat_message) + + # Stop the code executor. + await code_executor.stop() + + + asyncio.run(run_code_executor_agent()) + + In the following example, we show how to setup `CodeExecutorAgent` without `model_client` parameter for executing code blocks generated by other agents in a group chat using :py:class:`~autogen_ext.code_executors.docker.DockerCommandLineCodeExecutor` + + .. code-block:: python + + import asyncio + + from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor + from autogen_ext.models.openai import OpenAIChatCompletionClient + + from autogen_agentchat.agents import AssistantAgent, CodeExecutorAgent, ApprovalRequest, ApprovalResponse + from autogen_agentchat.conditions import MaxMessageTermination + from autogen_agentchat.teams import RoundRobinGroupChat + from autogen_agentchat.ui import Console + + termination_condition = MaxMessageTermination(3) + + + def group_chat_approval_func(request: ApprovalRequest) -> ApprovalResponse: + \"\"\"Approval function for group chat that allows basic Python operations.\"\"\" + # Allow common safe operations + safe_operations = ["print(", "import ", "def ", "class ", "if ", "for ", "while "] + if any(op in request.code for op in safe_operations): + return ApprovalResponse(approved=True, reason='Safe Python operation') + + # Deny file system operations in group chat + dangerous_operations = ["open(", "file(", "os.", "subprocess", "eval(", "exec("] + if any(op in request.code for op in dangerous_operations): + return ApprovalResponse(approved=False, reason='File system or dangerous operation not allowed') + + return ApprovalResponse(approved=True, reason='Operation approved') + + + async def main() -> None: + model_client = OpenAIChatCompletionClient(model="gpt-4o") + + # define the Docker CLI Code Executor + code_executor = DockerCommandLineCodeExecutor(work_dir="coding") + + # start the execution container + await code_executor.start() + + code_executor_agent = CodeExecutorAgent( + "code_executor_agent", + code_executor=code_executor, + approval_func=group_chat_approval_func + ) + coder_agent = AssistantAgent("coder_agent", model_client=model_client) + + groupchat = RoundRobinGroupChat( + participants=[coder_agent, code_executor_agent], termination_condition=termination_condition + ) + + task = "Write python code to print Hello World!" + await Console(groupchat.run_stream(task=task)) + + # stop the execution container + await code_executor.stop() + + + asyncio.run(main()) + + .. code-block:: text + + ---------- user ---------- + Write python code to print Hello World! + ---------- coder_agent ---------- + Certainly! Here's a simple Python code to print "Hello World!": + + ```python + print("Hello World!") + ``` + + You can run this code in any Python environment to display the message. + ---------- code_executor_agent ---------- + Hello World! + + In the following example, we show how to setup `CodeExecutorAgent` with `model_client` + that can generate its own code without the help of any other agent and executing it in + :py:class:`~autogen_ext.code_executors.docker.DockerCommandLineCodeExecutor`. + It also demonstrates using a model-based approval function that reviews the code for safety before execution. + + .. code-block:: python + + import asyncio + + from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_core.models import SystemMessage, UserMessage + + from autogen_agentchat.agents import CodeExecutorAgent, ApprovalRequest, ApprovalResponse + from autogen_agentchat.conditions import TextMessageTermination + from autogen_agentchat.ui import Console + + termination_condition = TextMessageTermination("code_executor_agent") + + + async def main() -> None: + model_client = OpenAIChatCompletionClient(model="gpt-4o") + + async def model_client_approval_func(request: ApprovalRequest) -> ApprovalResponse: + instruction = "Approve or reject the code in the last message based on whether it is dangerous or not. Use the following JSON format for your response: {approved: true/false, reason: 'your reason here'}" + response = await model_client.create( + messages=[SystemMessage(content=instruction)] + + request.context + + [UserMessage(content=request.code, source="user")], + json_output=ApprovalResponse, + ) + assert isinstance(response.content, str) + return ApprovalResponse.model_validate_json(response.content) + + # define the Docker CLI Code Executor + code_executor = DockerCommandLineCodeExecutor(work_dir="coding") + + # start the execution container + await code_executor.start() + + code_executor_agent = CodeExecutorAgent( + "code_executor_agent", + code_executor=code_executor, + model_client=model_client, + approval_func=model_client_approval_func, + ) + + task = "Write python code to print Hello World!" + await Console(code_executor_agent.run_stream(task=task)) + + # stop the execution container + await code_executor.stop() + + + asyncio.run(main()) + + + .. code-block:: text + + ---------- user ---------- + Write python code to print Hello World! + ---------- code_executor_agent ---------- + Certainly! Here is a simple Python code to print "Hello World!" to the console: + + ```python + print("Hello World!") + ``` + + Let's execute it to confirm the output. + ---------- code_executor_agent ---------- + Hello World! + + ---------- code_executor_agent ---------- + The code has been executed successfully, and it printed "Hello World!" as expected. If you have any more requests or questions, feel free to ask! + """ + DEFAULT_TERMINAL_DESCRIPTION = "A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks)." + DEFAULT_AGENT_DESCRIPTION = "A Code Execution Agent that generates and executes Python and shell scripts based on user instructions. It ensures correctness, efficiency, and minimal errors while gracefully handling edge cases." + DEFAULT_SYSTEM_MESSAGE = "You are a Code Execution Agent. Your role is to generate and execute Python code and shell scripts based on user instructions, ensuring correctness, efficiency, and minimal errors. Handle edge cases gracefully. Python code should be provided in ```python code blocks, and sh shell scripts should be provided in ```sh code blocks for execution." + NO_CODE_BLOCKS_FOUND_MESSAGE = "No code blocks found in the thread. Please provide at least one markdown-encoded code block to execute (i.e., quoting code in ```python or ```sh code blocks)." + DEFAULT_SUPPORTED_LANGUAGES = ["python", "sh"] + component_config_schema = CodeExecutorAgentConfig component_provider_override = "autogen_agentchat.agents.CodeExecutorAgent" @@ -109,52 +430,296 @@ def __init__( name: str, code_executor: CodeExecutor, *, - description: str = "A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks).", + model_client: ChatCompletionClient | None = None, + model_context: ChatCompletionContext | None = None, + model_client_stream: bool = False, + max_retries_on_error: int = 0, + description: str | None = None, + system_message: str | None = DEFAULT_SYSTEM_MESSAGE, sources: Sequence[str] | None = None, + supported_languages: List[str] | None = None, + approval_func: Optional[ApprovalFuncType] = None, ) -> None: + if description is None: + if model_client is None: + description = CodeExecutorAgent.DEFAULT_TERMINAL_DESCRIPTION + else: + description = CodeExecutorAgent.DEFAULT_AGENT_DESCRIPTION + super().__init__(name=name, description=description) self._code_executor = code_executor self._sources = sources + self._model_client_stream = model_client_stream + self._max_retries_on_error = max_retries_on_error + self._approval_func = approval_func + self._approval_func_is_async = approval_func is not None and iscoroutinefunction(approval_func) + + if supported_languages is not None: + self._supported_languages = supported_languages + else: + self._supported_languages = CodeExecutorAgent.DEFAULT_SUPPORTED_LANGUAGES + + self._supported_languages_regex = "|".join(re.escape(lang) for lang in self._supported_languages) + + self._model_client = None + if model_client is not None: + self._model_client = model_client + + if model_context is not None: + self._model_context = model_context + else: + self._model_context = UnboundedChatCompletionContext() + + self._system_messaages: List[SystemMessage] = [] + if system_message is None: + self._system_messages = [] + else: + self._system_messages = [SystemMessage(content=system_message)] + + if self._max_retries_on_error > 0: + if not self._model_client or not self._model_client.model_info: + raise ValueError("model_client.model_info must be provided when max_retries_on_error > 0") + if not self._model_client.model_info["structured_output"]: + raise ValueError("Specified model_client doesn't support structured output mode.") @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: """The types of messages that the code executor agent produces.""" return (TextMessage,) - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + @property + def model_context(self) -> ChatCompletionContext: + """ + The model context in use by the agent. + """ + return self._model_context + + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: + async for message in self.on_messages_stream(messages, cancellation_token): + if isinstance(message, Response): + return message + raise AssertionError("The stream should have returned the final result.") + + async def on_messages_stream( + self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken + ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]: + """ + Process the incoming messages with the assistant agent and yield events/responses as they happen. + """ + + # Gather all relevant state here + agent_name = self.name + model_context = self._model_context + system_messages = self._system_messages + model_client = self._model_client + model_client_stream = self._model_client_stream + max_retries_on_error = self._max_retries_on_error + + execution_result: CodeResult | None = None + if model_client is None: # default behaviour for backward compatibility + # execute generated code if present + code_blocks: List[CodeBlock] = await self.extract_code_blocks_from_messages(messages) + if not code_blocks: + yield Response( + chat_message=TextMessage( + content=self.NO_CODE_BLOCKS_FOUND_MESSAGE, + source=agent_name, + ) + ) + return + execution_result = await self.execute_code_block(code_blocks, cancellation_token) + yield Response(chat_message=TextMessage(content=execution_result.output, source=self.name)) + return + + inner_messages: List[BaseAgentEvent | BaseChatMessage] = [] + + for nth_try in range(max_retries_on_error + 1): # Do one default generation, execution and inference loop + # Step 1: Add new user/handoff messages to the model context + await self._add_messages_to_context( + model_context=model_context, + messages=messages, + ) + + # Step 2: Run inference with the model context + model_result = None + async for inference_output in self._call_llm( + model_client=model_client, + model_client_stream=model_client_stream, + system_messages=system_messages, + model_context=model_context, + agent_name=agent_name, + cancellation_token=cancellation_token, + ): + if isinstance(inference_output, CreateResult): + model_result = inference_output + else: + # Streaming chunk event + yield inference_output + + assert model_result is not None, "No model result was produced." + + # Step 3: [NEW] If the model produced a hidden "thought," yield it as an event + if model_result.thought: + thought_event = ThoughtEvent(content=model_result.thought, source=agent_name) + yield thought_event + inner_messages.append(thought_event) + + # Step 4: Add the assistant message to the model context (including thought if present) + await model_context.add_message( + AssistantMessage( + content=model_result.content, + source=agent_name, + thought=getattr(model_result, "thought", None), + ) + ) + + # Step 5: Extract the code blocks from inferred text + assert isinstance(model_result.content, str), "Expected inferred model_result.content to be of type str." + code_blocks = self._extract_markdown_code_blocks(str(model_result.content)) + + # Step 6: Exit the loop if no code blocks found + if not code_blocks: + yield Response( + chat_message=TextMessage( + content=str(model_result.content), + source=agent_name, + ) + ) + return + + # Step 7: Yield a CodeGenerationEvent + inferred_text_message: CodeGenerationEvent = CodeGenerationEvent( + retry_attempt=nth_try, + content=model_result.content, + code_blocks=code_blocks, + source=agent_name, + ) + + yield inferred_text_message + + # Step 8: Execute the extracted code blocks + execution_result = await self.execute_code_block(inferred_text_message.code_blocks, cancellation_token) + + # Step 9: Update model context with the code execution result + await model_context.add_message( + UserMessage( + content=execution_result.output, + source=agent_name, + ) + ) + + # Step 10: Yield a CodeExecutionEvent + yield CodeExecutionEvent(retry_attempt=nth_try, result=execution_result, source=self.name) + + # If execution was successful or last retry, then exit + if execution_result.exit_code == 0 or nth_try == max_retries_on_error: + break + + # Step 11: If exit code is non-zero and retries are available then + # make an inference asking if we should retry or not + chat_context = await model_context.get_messages() + + retry_prompt = ( + f"The most recent code execution resulted in an error:\n{execution_result.output}\n\n" + "Should we attempt to resolve it? Please respond with:\n" + "- A boolean value for 'retry' indicating whether it should be retried.\n" + "- A detailed explanation in 'reason' that identifies the issue, justifies your decision to retry or not, and outlines how you would resolve the error if a retry is attempted." + ) + + chat_context = chat_context + [ + UserMessage( + content=retry_prompt, + source=agent_name, + ) + ] + + response = await model_client.create(messages=chat_context, json_output=RetryDecision) + + assert isinstance( + response.content, str + ), "Expected structured response for retry decision to be of type str." + should_retry_generation = RetryDecision.model_validate_json(str(response.content)) + + # Exit if no-retry is needed + if not should_retry_generation.retry: + break + + yield CodeGenerationEvent( + retry_attempt=nth_try, + content=f"Attempt number: {nth_try + 1}\nProposed correction: {should_retry_generation.reason}", + code_blocks=[], + source=agent_name, + ) + + # Always reflect on the execution result + async for reflection_response in CodeExecutorAgent._reflect_on_code_block_results_flow( + system_messages=system_messages, + model_client=model_client, + model_client_stream=model_client_stream, + model_context=model_context, + agent_name=agent_name, + inner_messages=inner_messages, + ): + yield reflection_response # Last reflection_response is of type Response so it will finish the routine + + async def extract_code_blocks_from_messages(self, messages: Sequence[BaseChatMessage]) -> List[CodeBlock]: # Extract code blocks from the messages. code_blocks: List[CodeBlock] = [] for msg in messages: - if isinstance(msg, TextMessage): - if self._sources is None or msg.source in self._sources: + if self._sources is None or msg.source in self._sources: + if isinstance(msg, TextMessage): code_blocks.extend(self._extract_markdown_code_blocks(msg.content)) - if code_blocks: - # Execute the code blocks. - result = await self._code_executor.execute_code_blocks(code_blocks, cancellation_token=cancellation_token) - - code_output = result.output - if code_output.strip() == "": - # No output - code_output = f"The script ran but produced no output to console. The POSIX exit code was: {result.exit_code}. If you were expecting output, consider revising the script to ensure content is printed to stdout." - elif result.exit_code != 0: - # Error - code_output = f"The script ran, then exited with an error (POSIX exit code: {result.exit_code})\nIts output was:\n{result.output}" - - return Response(chat_message=TextMessage(content=code_output, source=self.name)) - else: - return Response( - chat_message=TextMessage( - content="No code blocks found in the thread. Please provide at least one markdown-encoded code block to execute (i.e., quoting code in ```python or ```sh code blocks).", - source=self.name, + # TODO: handle other message types if needed + return code_blocks + + async def execute_code_block( + self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken + ) -> CodeResult: + # Check for approval before executing code blocks + if self._approval_func is not None: + # Combine all code blocks into a single string for approval + combined_code = "\n\n".join([f"```{block.language}\n{block.code}\n```" for block in code_blocks]) + + # Get the current context from model_context + context_messages = await self._model_context.get_messages() + + # Create approval request + approval_request = ApprovalRequest(code=combined_code, context=context_messages) + + # Get approval (handle both sync and async functions) + if self._approval_func_is_async: + # Cast to AsyncApprovalFunc for proper typing + async_func = cast(AsyncApprovalFunc, self._approval_func) + approval_response = await async_func(approval_request) + else: + # Cast to SyncApprovalFunc for proper typing + sync_func = cast(SyncApprovalFunc, self._approval_func) + approval_response = sync_func(approval_request) + + # If not approved, return error result + if not approval_response.approved: + return CodeResult( + exit_code=1, output=f"Code execution was not approved. Reason: {approval_response.reason}" ) - ) + + # Execute the code blocks. + result = await self._code_executor.execute_code_blocks(code_blocks, cancellation_token=cancellation_token) + + if result.output.strip() == "": + # No output + result.output = f"The script ran but produced no output to console. The POSIX exit code was: {result.exit_code}. If you were expecting output, consider revising the script to ensure content is printed to stdout." + elif result.exit_code != 0: + # Error + result.output = f"The script ran, then exited with an error (POSIX exit code: {result.exit_code})\nIts output was:\n{result.output}" + + return result async def on_reset(self, cancellation_token: CancellationToken) -> None: - """It it's a no-op as the code executor agent has no mutable state.""" + """Its a no-op as the code executor agent has no mutable state.""" pass def _extract_markdown_code_blocks(self, markdown_text: str) -> List[CodeBlock]: - pattern = re.compile(r"```(?:\s*([\w\+\-]+))?\n([\s\S]*?)```") + pattern = re.compile(rf"```(?:\s*({self._supported_languages_regex}))\n([\s\S]*?)```", re.IGNORECASE) matches = pattern.findall(markdown_text) code_blocks: List[CodeBlock] = [] for match in matches: @@ -164,18 +729,153 @@ def _extract_markdown_code_blocks(self, markdown_text: str) -> List[CodeBlock]: return code_blocks def _to_config(self) -> CodeExecutorAgentConfig: + if self._approval_func is not None: + raise ValueError( + "Cannot serialize CodeExecutorAgent with approval_func set. The approval function is not serializable." + ) + return CodeExecutorAgentConfig( name=self.name, + model_client=(self._model_client.dump_component() if self._model_client is not None else None), code_executor=self._code_executor.dump_component(), description=self.description, sources=list(self._sources) if self._sources is not None else None, + system_message=( + self._system_messages[0].content + if self._system_messages and isinstance(self._system_messages[0].content, str) + else None + ), + model_client_stream=self._model_client_stream, + model_context=self._model_context.dump_component(), + supported_languages=self._supported_languages, ) @classmethod def _from_config(cls, config: CodeExecutorAgentConfig) -> Self: return cls( name=config.name, + model_client=( + ChatCompletionClient.load_component(config.model_client) if config.model_client is not None else None + ), code_executor=CodeExecutor.load_component(config.code_executor), description=config.description, sources=config.sources, + system_message=config.system_message, + model_client_stream=config.model_client_stream, + model_context=ChatCompletionContext.load_component(config.model_context) if config.model_context else None, + supported_languages=config.supported_languages, + approval_func=None, # approval_func cannot be serialized, so it's always None when loading from config + ) + + @staticmethod + def _get_compatible_context(model_client: ChatCompletionClient, messages: List[LLMMessage]) -> Sequence[LLMMessage]: + """Ensure that the messages are compatible with the underlying client, by removing images if needed.""" + if model_client.model_info["vision"]: + return messages + else: + return remove_images(messages) + + @classmethod + async def _call_llm( + cls, + model_client: ChatCompletionClient, + model_client_stream: bool, + system_messages: List[SystemMessage], + model_context: ChatCompletionContext, + agent_name: str, + cancellation_token: CancellationToken, + ) -> AsyncGenerator[Union[CreateResult, ModelClientStreamingChunkEvent], None]: + """ + Perform a model inference and yield either streaming chunk events or the final CreateResult. + """ + all_messages = await model_context.get_messages() + llm_messages = cls._get_compatible_context(model_client=model_client, messages=system_messages + all_messages) + + if model_client_stream: + model_result: Optional[CreateResult] = None + async for chunk in model_client.create_stream( + llm_messages, tools=[], cancellation_token=cancellation_token + ): + if isinstance(chunk, CreateResult): + model_result = chunk + elif isinstance(chunk, str): + yield ModelClientStreamingChunkEvent(content=chunk, source=agent_name) + else: + raise RuntimeError(f"Invalid chunk type: {type(chunk)}") + if model_result is None: + raise RuntimeError("No final model result in streaming mode.") + yield model_result + else: + model_result = await model_client.create(llm_messages, tools=[], cancellation_token=cancellation_token) + yield model_result + + @staticmethod + async def _add_messages_to_context( + model_context: ChatCompletionContext, + messages: Sequence[BaseChatMessage], + ) -> None: + """ + Add incoming messages to the model context. + """ + for msg in messages: + if isinstance(msg, HandoffMessage): + for llm_msg in msg.context: + await model_context.add_message(llm_msg) + await model_context.add_message(msg.to_model_message()) + + @classmethod + async def _reflect_on_code_block_results_flow( + cls, + system_messages: List[SystemMessage], + model_client: ChatCompletionClient, + model_client_stream: bool, + model_context: ChatCompletionContext, + agent_name: str, + inner_messages: List[BaseAgentEvent | BaseChatMessage], + ) -> AsyncGenerator[Response | ModelClientStreamingChunkEvent | ThoughtEvent, None]: + """ + If reflect_on_code_block_results=True, we do another inference based on tool results + and yield the final text response (or streaming chunks). + """ + all_messages = system_messages + await model_context.get_messages() + llm_messages = cls._get_compatible_context(model_client=model_client, messages=all_messages) + + reflection_result: Optional[CreateResult] = None + + if model_client_stream: + async for chunk in model_client.create_stream(llm_messages): + if isinstance(chunk, CreateResult): + reflection_result = chunk + elif isinstance(chunk, str): + yield ModelClientStreamingChunkEvent(content=chunk, source=agent_name) + else: + raise RuntimeError(f"Invalid chunk type: {type(chunk)}") + else: + reflection_result = await model_client.create(llm_messages) + + if not reflection_result or not isinstance(reflection_result.content, str): + raise RuntimeError("Reflect on tool use produced no valid text response.") + + # --- NEW: If the reflection produced a thought, yield it --- + if reflection_result.thought: + thought_event = ThoughtEvent(content=reflection_result.thought, source=agent_name) + yield thought_event + inner_messages.append(thought_event) + + # Add to context (including thought if present) + await model_context.add_message( + AssistantMessage( + content=reflection_result.content, + source=agent_name, + thought=getattr(reflection_result, "thought", None), + ) + ) + + yield Response( + chat_message=TextMessage( + content=reflection_result.content, + source=agent_name, + models_usage=reflection_result.usage, + ), + inner_messages=inner_messages, ) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_message_filter_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_message_filter_agent.py new file mode 100644 index 000000000000..0905e694d513 --- /dev/null +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_message_filter_agent.py @@ -0,0 +1,203 @@ +from typing import AsyncGenerator, List, Literal, Optional, Sequence, Union + +from autogen_core import CancellationToken, Component, ComponentModel +from pydantic import BaseModel + +from autogen_agentchat.agents import BaseChatAgent +from autogen_agentchat.base import Response +from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage + +# ------------------------------ +# Message Filter Config +# ------------------------------ + + +class PerSourceFilter(BaseModel): + source: str + position: Optional[Literal["first", "last"]] = None + count: Optional[int] = None + + +class MessageFilterConfig(BaseModel): + per_source: List[PerSourceFilter] + + +# ------------------------------ +# Component Config +# ------------------------------ + + +class MessageFilterAgentConfig(BaseModel): + name: str + wrapped_agent: ComponentModel + filter: MessageFilterConfig + + +# ------------------------------ +# Message Filter Agent +# ------------------------------ + + +class MessageFilterAgent(BaseChatAgent, Component[MessageFilterAgentConfig]): + """ + A wrapper agent that filters incoming messages before passing them to the inner agent. + + .. warning:: + + This is an experimental feature, and the API will change in the future releases. + + This is useful in scenarios like multi-agent workflows where an agent should only + process a subset of the full message history—for example, only the last message + from each upstream agent, or only the first message from a specific source. + + Filtering is configured using :class:`MessageFilterConfig`, which supports: + - Filtering by message source (e.g., only messages from "user" or another agent) + - Selecting the first N or last N messages from each source + - If position is `None`, all messages from that source are included + + This agent is compatible with both direct message passing and team-based execution + such as :class:`~autogen_agentchat.teams.GraphFlow`. + + Example: + >>> agent_a = MessageFilterAgent( + ... name="A", + ... wrapped_agent=some_other_agent, + ... filter=MessageFilterConfig( + ... per_source=[ + ... PerSourceFilter(source="user", position="first", count=1), + ... PerSourceFilter(source="B", position="last", count=2), + ... ] + ... ), + ... ) + + Example use case with Graph: + Suppose you have a looping multi-agent graph: A → B → A → B → C. + + You want: + - A to only see the user message and the last message from B + - B to see the user message, last message from A, and its own prior responses (for reflection) + - C to see the user message and the last message from B + + Wrap the agents like so: + + >>> agent_a = MessageFilterAgent( + ... name="A", + ... wrapped_agent=agent_a_inner, + ... filter=MessageFilterConfig( + ... per_source=[ + ... PerSourceFilter(source="user", position="first", count=1), + ... PerSourceFilter(source="B", position="last", count=1), + ... ] + ... ), + ... ) + + >>> agent_b = MessageFilterAgent( + ... name="B", + ... wrapped_agent=agent_b_inner, + ... filter=MessageFilterConfig( + ... per_source=[ + ... PerSourceFilter(source="user", position="first", count=1), + ... PerSourceFilter(source="A", position="last", count=1), + ... PerSourceFilter(source="B", position="last", count=10), + ... ] + ... ), + ... ) + + >>> agent_c = MessageFilterAgent( + ... name="C", + ... wrapped_agent=agent_c_inner, + ... filter=MessageFilterConfig( + ... per_source=[ + ... PerSourceFilter(source="user", position="first", count=1), + ... PerSourceFilter(source="B", position="last", count=1), + ... ] + ... ), + ... ) + + Then define the graph: + + >>> graph = DiGraph( + ... nodes={ + ... "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + ... "B": DiGraphNode( + ... name="B", + ... edges=[ + ... DiGraphEdge(target="C", condition="exit"), + ... DiGraphEdge(target="A", condition="loop"), + ... ], + ... ), + ... "C": DiGraphNode(name="C", edges=[]), + ... }, + ... default_start_node="A", + ... ) + + This will ensure each agent sees only what is needed for its decision or action logic. + """ + + component_config_schema = MessageFilterAgentConfig + component_provider_override = "autogen_agentchat.agents.MessageFilterAgent" + + def __init__( + self, + name: str, + wrapped_agent: BaseChatAgent, + filter: MessageFilterConfig, + ): + super().__init__(name=name, description=f"{wrapped_agent.description} (with message filtering)") + self._wrapped_agent = wrapped_agent + self._filter = filter + + @property + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: + return self._wrapped_agent.produced_message_types + + def _apply_filter(self, messages: Sequence[BaseChatMessage]) -> Sequence[BaseChatMessage]: + result: List[BaseChatMessage] = [] + + for source_filter in self._filter.per_source: + msgs = [m for m in messages if m.source == source_filter.source] + + if source_filter.position == "first" and source_filter.count: + msgs = msgs[: source_filter.count] + elif source_filter.position == "last" and source_filter.count: + msgs = msgs[-source_filter.count :] + + result.extend(msgs) + + return result + + async def on_messages( + self, + messages: Sequence[BaseChatMessage], + cancellation_token: CancellationToken, + ) -> Response: + filtered = self._apply_filter(messages) + return await self._wrapped_agent.on_messages(filtered, cancellation_token) + + async def on_messages_stream( + self, + messages: Sequence[BaseChatMessage], + cancellation_token: CancellationToken, + ) -> AsyncGenerator[Union[BaseAgentEvent, BaseChatMessage, Response], None]: + filtered = self._apply_filter(messages) + async for item in self._wrapped_agent.on_messages_stream(filtered, cancellation_token): + yield item + + async def on_reset(self, cancellation_token: CancellationToken) -> None: + await self._wrapped_agent.on_reset(cancellation_token) + + def _to_config(self) -> MessageFilterAgentConfig: + return MessageFilterAgentConfig( + name=self.name, + wrapped_agent=self._wrapped_agent.dump_component(), + filter=self._filter, + ) + + @classmethod + def _from_config(cls, config: MessageFilterAgentConfig) -> "MessageFilterAgent": + wrapped = BaseChatAgent.load_component(config.wrapped_agent) + return cls( + name=config.name, + wrapped_agent=wrapped, + filter=config.filter, + ) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py index 2eba918714b7..31f26db65ce1 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py @@ -1,6 +1,10 @@ from typing import Any, AsyncGenerator, List, Mapping, Sequence from autogen_core import CancellationToken, Component, ComponentModel +from autogen_core.model_context import ( + ChatCompletionContext, + UnboundedChatCompletionContext, +) from autogen_core.models import ChatCompletionClient, LLMMessage, SystemMessage, UserMessage from pydantic import BaseModel from typing_extensions import Self @@ -10,9 +14,9 @@ from ..base import TaskResult, Team from ..messages import ( - AgentEvent, + BaseAgentEvent, BaseChatMessage, - ChatMessage, + HandoffMessage, ModelClientStreamingChunkEvent, TextMessage, ) @@ -28,6 +32,7 @@ class SocietyOfMindAgentConfig(BaseModel): description: str | None = None instruction: str | None = None response_prompt: str | None = None + model_context: ComponentModel | None = None class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]): @@ -39,6 +44,16 @@ class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]): Once the response is generated, the agent resets the inner team by calling :meth:`Team.reset`. + Limit context size sent to the model: + + You can limit the number of messages sent to the model by setting + the `model_context` parameter to a :class:`~autogen_core.model_context.BufferedChatCompletionContext`. + This will limit the number of recent messages sent to the model and can be useful + when the model has a limit on the number of tokens it can process. + You can also create your own model context by subclassing + :class:`~autogen_core.model_context.ChatCompletionContext`. + + Args: name (str): The name of the agent. team (Team): The team of agents to use. @@ -48,6 +63,8 @@ class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]): Defaults to :attr:`DEFAULT_INSTRUCTION`. It assumes the role of 'system'. response_prompt (str, optional): The response prompt to use when generating a response using the inner team's messages. Defaults to :attr:`DEFAULT_RESPONSE_PROMPT`. It assumes the role of 'system'. + model_context (ChatCompletionContext | None, optional): The model context for storing and retrieving :class:`~autogen_core.models.LLMMessage`. It can be preloaded with initial messages. The initial messages will be cleared when the agent is reset. + Example: @@ -115,6 +132,7 @@ def __init__( description: str = DEFAULT_DESCRIPTION, instruction: str = DEFAULT_INSTRUCTION, response_prompt: str = DEFAULT_RESPONSE_PROMPT, + model_context: ChatCompletionContext | None = None, ) -> None: super().__init__(name=name, description=description) self._team = team @@ -122,11 +140,23 @@ def __init__( self._instruction = instruction self._response_prompt = response_prompt + if model_context is not None: + self._model_context = model_context + else: + self._model_context = UnboundedChatCompletionContext() + @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: return (TextMessage,) - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + @property + def model_context(self) -> ChatCompletionContext: + """ + The model context in use by the agent. + """ + return self._model_context + + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: # Call the stream method and collect the messages. response: Response | None = None async for msg in self.on_messages_stream(messages, cancellation_token): @@ -136,23 +166,38 @@ async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: return response async def on_messages_stream( - self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken - ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]: + self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken + ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]: # Prepare the task for the team of agents. - task = list(messages) + task_messages = list(messages) # Run the team of agents. result: TaskResult | None = None - inner_messages: List[AgentEvent | ChatMessage] = [] - count = 0 - async for inner_msg in self._team.run_stream(task=task, cancellation_token=cancellation_token): + inner_messages: List[BaseAgentEvent | BaseChatMessage] = [] + model_context = self._model_context + + prev_content = await model_context.get_messages() + if len(prev_content) > 0: + prev_message = HandoffMessage( + content="relevant previous messages", + source=self.name, + target="", + context=prev_content, + ) + task_messages = [prev_message] + task_messages + + if len(task_messages) == 0: + task = None + else: + task = task_messages + + # Use the new output_task_messages parameter to avoid fragile count-based logic + async for inner_msg in self._team.run_stream( + task=task, cancellation_token=cancellation_token, output_task_messages=False + ): if isinstance(inner_msg, TaskResult): result = inner_msg else: - count += 1 - if count <= len(task): - # Skip the task messages. - continue yield inner_msg if isinstance(inner_msg, ModelClientStreamingChunkEvent): # Skip the model client streaming chunk events. @@ -162,31 +207,65 @@ async def on_messages_stream( if len(inner_messages) == 0: yield Response( - chat_message=TextMessage(source=self.name, content="No response."), inner_messages=inner_messages + chat_message=TextMessage(source=self.name, content="No response."), + inner_messages=[], + # Response's inner_messages should be empty. Cause that mean is response to outer world. ) else: + llm_messages: List[LLMMessage] = [] + + if self._model_client.model_info.get("multiple_system_messages", False): + # The model client supports multiple system messages, so we + llm_messages.append(SystemMessage(content=self._instruction)) + else: + # The model client does not support multiple system messages, so we + llm_messages.append(UserMessage(content=self._instruction, source="user")) + # Generate a response using the model client. - llm_messages: List[LLMMessage] = [SystemMessage(content=self._instruction)] - llm_messages.extend( - [ - UserMessage(content=message.content, source=message.source) - for message in inner_messages - if isinstance(message, BaseChatMessage) - ] - ) - llm_messages.append(SystemMessage(content=self._response_prompt)) + for message in inner_messages: + if isinstance(message, BaseChatMessage): + llm_messages.append(message.to_model_message()) + + if self._model_client.model_info.get("multiple_system_messages", False): + # The model client supports multiple system messages, so we + llm_messages.append(SystemMessage(content=self._response_prompt)) + else: + # The model client does not support multiple system messages, so we + llm_messages.append(UserMessage(content=self._response_prompt, source="user")) completion = await self._model_client.create(messages=llm_messages, cancellation_token=cancellation_token) assert isinstance(completion.content, str) yield Response( chat_message=TextMessage(source=self.name, content=completion.content, models_usage=completion.usage), - inner_messages=inner_messages, + inner_messages=[], + # Response's inner_messages should be empty. Cause that mean is response to outer world. ) + # Add new user/handoff messages to the model context + await self._add_messages_to_context( + model_context=model_context, + messages=messages, + ) + # Reset the team. await self._team.reset() + @staticmethod + async def _add_messages_to_context( + model_context: ChatCompletionContext, + messages: Sequence[BaseChatMessage], + ) -> None: + """ + Add incoming messages to the model context. + """ + for msg in messages: + if isinstance(msg, HandoffMessage): + for llm_msg in msg.context: + await model_context.add_message(llm_msg) + await model_context.add_message(msg.to_model_message()) + async def on_reset(self, cancellation_token: CancellationToken) -> None: await self._team.reset() + await self._model_context.clear() async def save_state(self) -> Mapping[str, Any]: team_state = await self._team.save_state() @@ -205,6 +284,7 @@ def _to_config(self) -> SocietyOfMindAgentConfig: description=self.description, instruction=self._instruction, response_prompt=self._response_prompt, + model_context=self._model_context.dump_component(), ) @classmethod @@ -218,4 +298,5 @@ def _from_config(cls, config: SocietyOfMindAgentConfig) -> Self: description=config.description or cls.DEFAULT_DESCRIPTION, instruction=config.instruction or cls.DEFAULT_INSTRUCTION, response_prompt=config.response_prompt or cls.DEFAULT_RESPONSE_PROMPT, + model_context=ChatCompletionContext.load_component(config.model_context) if config.model_context else None, ) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_user_proxy_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_user_proxy_agent.py index 3ca0ec890324..af78f64c93c8 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_user_proxy_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_user_proxy_agent.py @@ -10,7 +10,7 @@ from typing_extensions import Self from ..base import Response -from ..messages import AgentEvent, ChatMessage, HandoffMessage, TextMessage, UserInputRequestedEvent +from ..messages import BaseAgentEvent, BaseChatMessage, HandoffMessage, TextMessage, UserInputRequestedEvent from ._base_chat_agent import BaseChatAgent SyncInputFunc = Callable[[str], str] @@ -82,6 +82,7 @@ async def simple_user_agent(): cancellation_token=CancellationToken(), ) ) + assert isinstance(response.chat_message, TextMessage) print(f"Your name is {response.chat_message.content}") Example: @@ -117,6 +118,7 @@ async def cancellable_user_agent(): ) ) response = await agent_task + assert isinstance(response.chat_message, TextMessage) print(f"Your name is {response.chat_message.content}") except Exception as e: print(f"Exception: {e}") @@ -168,11 +170,11 @@ def __init__( self._is_async = iscoroutinefunction(self.input_func) @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: """Message types this agent can produce.""" return (TextMessage, HandoffMessage) - def _get_latest_handoff(self, messages: Sequence[ChatMessage]) -> Optional[HandoffMessage]: + def _get_latest_handoff(self, messages: Sequence[BaseChatMessage]) -> Optional[HandoffMessage]: """Find the HandoffMessage in the message sequence that addresses this agent.""" if len(messages) > 0 and isinstance(messages[-1], HandoffMessage): if messages[-1].target == self.name: @@ -199,15 +201,15 @@ async def _get_input(self, prompt: str, cancellation_token: Optional[Cancellatio except Exception as e: raise RuntimeError(f"Failed to get user input: {str(e)}") from e - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: async for message in self.on_messages_stream(messages, cancellation_token): if isinstance(message, Response): return message raise AssertionError("The stream should have returned the final result.") async def on_messages_stream( - self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken - ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]: + self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken + ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]: """Handle incoming messages by requesting user input.""" try: # Check for handoff first diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_chat_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/base/_chat_agent.py index ec885ee7f8cb..5bc8e803844a 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_chat_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/base/_chat_agent.py @@ -3,9 +3,9 @@ from typing import Any, AsyncGenerator, Mapping, Sequence from autogen_core import CancellationToken, ComponentBase -from pydantic import BaseModel +from pydantic import BaseModel, SerializeAsAny -from ..messages import AgentEvent, ChatMessage +from ..messages import BaseAgentEvent, BaseChatMessage from ._task import TaskRunner @@ -13,12 +13,12 @@ class Response: """A response from calling :meth:`ChatAgent.on_messages`.""" - chat_message: ChatMessage + chat_message: SerializeAsAny[BaseChatMessage] """A chat message produced by the agent as the response.""" - inner_messages: Sequence[AgentEvent | ChatMessage] | None = None - """Inner messages produced by the agent, they can be :class:`AgentEvent` - or :class:`ChatMessage`.""" + inner_messages: Sequence[SerializeAsAny[BaseAgentEvent | BaseChatMessage]] | None = None + """Inner messages produced by the agent, they can be :class:`BaseAgentEvent` + or :class:`BaseChatMessage`.""" class ChatAgent(ABC, TaskRunner, ComponentBase[BaseModel]): @@ -43,20 +43,20 @@ def description(self) -> str: @property @abstractmethod - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: """The types of messages that the agent produces in the - :attr:`Response.chat_message` field. They must be :class:`ChatMessage` types.""" + :attr:`Response.chat_message` field. They must be :class:`BaseChatMessage` types.""" ... @abstractmethod - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: """Handles incoming messages and returns a response.""" ... @abstractmethod def on_messages_stream( - self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken - ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]: + self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken + ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]: """Handles incoming messages and returns a stream of inner messages and and the final item is the response.""" ... diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_task.py b/python/packages/autogen-agentchat/src/autogen_agentchat/base/_task.py index 90e319ee36c6..b858b4a4517c 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_task.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/base/_task.py @@ -1,16 +1,15 @@ -from dataclasses import dataclass from typing import AsyncGenerator, Protocol, Sequence from autogen_core import CancellationToken +from pydantic import BaseModel, SerializeAsAny -from ..messages import AgentEvent, ChatMessage +from ..messages import BaseAgentEvent, BaseChatMessage -@dataclass -class TaskResult: +class TaskResult(BaseModel): """Result of running a task.""" - messages: Sequence[AgentEvent | ChatMessage] + messages: Sequence[SerializeAsAny[BaseAgentEvent | BaseChatMessage]] """Messages produced by the task.""" stop_reason: str | None = None @@ -23,8 +22,9 @@ class TaskRunner(Protocol): async def run( self, *, - task: str | ChatMessage | Sequence[ChatMessage] | None = None, + task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None, cancellation_token: CancellationToken | None = None, + output_task_messages: bool = True, ) -> TaskResult: """Run the task and return the result. @@ -32,15 +32,22 @@ async def run( The runner is stateful and a subsequent call to this method will continue from where the previous call left off. If the task is not specified, - the runner will continue with the current task.""" + the runner will continue with the current task. + + Args: + task: The task to run. Can be a string, a single message, or a sequence of messages. + cancellation_token: The cancellation token to kill the task immediately. + output_task_messages: Whether to include task messages in :attr:`TaskResult.messages`. Defaults to True for backward compatibility. + """ ... def run_stream( self, *, - task: str | ChatMessage | Sequence[ChatMessage] | None = None, + task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None, cancellation_token: CancellationToken | None = None, - ) -> AsyncGenerator[AgentEvent | ChatMessage | TaskResult, None]: + output_task_messages: bool = True, + ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None]: """Run the task and produces a stream of messages and the final result :class:`TaskResult` as the last item in the stream. @@ -48,5 +55,11 @@ def run_stream( The runner is stateful and a subsequent call to this method will continue from where the previous call left off. If the task is not specified, - the runner will continue with the current task.""" + the runner will continue with the current task. + + Args: + task: The task to run. Can be a string, a single message, or a sequence of messages. + cancellation_token: The cancellation token to kill the task immediately. + output_task_messages: Whether to include task messages in the output stream. Defaults to True for backward compatibility. + """ ... diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_team.py b/python/packages/autogen-agentchat/src/autogen_agentchat/base/_team.py index 4cbdc9741c5c..e39aedaa67c8 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_team.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/base/_team.py @@ -10,6 +10,20 @@ class Team(ABC, TaskRunner, ComponentBase[BaseModel]): component_type = "team" + @property + @abstractmethod + def name(self) -> str: + """The name of the team. This is used by team to uniquely identify itself + in a larger team of teams.""" + ... + + @property + @abstractmethod + def description(self) -> str: + """A description of the team. This is used to provide context about the + team and its purpose to its parent orchestrator.""" + ... + @abstractmethod async def reset(self) -> None: """Reset the team and all its participants to its initial state.""" diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_termination.py b/python/packages/autogen-agentchat/src/autogen_agentchat/base/_termination.py index d8a3adb96818..5dd720c51619 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_termination.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/base/_termination.py @@ -6,7 +6,7 @@ from pydantic import BaseModel from typing_extensions import Self -from ..messages import AgentEvent, ChatMessage, StopMessage +from ..messages import BaseAgentEvent, BaseChatMessage, StopMessage class TerminatedException(BaseException): ... @@ -15,7 +15,7 @@ class TerminatedException(BaseException): ... class TerminationCondition(ABC, ComponentBase[BaseModel]): """A stateful condition that determines when a conversation should be terminated. - A termination condition is a callable that takes a sequence of ChatMessage objects + A termination condition is a callable that takes a sequence of BaseChatMessage objects since the last time the condition was called, and returns a StopMessage if the conversation should be terminated, or None otherwise. Once a termination condition has been reached, it must be reset before it can be used again. @@ -56,7 +56,7 @@ def terminated(self) -> bool: ... @abstractmethod - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: """Check if the conversation should be terminated based on the messages received since the last time the condition was called. Return a StopMessage if the conversation should be terminated, or None otherwise. @@ -102,7 +102,7 @@ def __init__(self, *conditions: TerminationCondition) -> None: def terminated(self) -> bool: return all(condition.terminated for condition in self._conditions) - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: if self.terminated: raise TerminatedException("Termination condition has already been reached.") # Check all remaining conditions. @@ -153,13 +153,14 @@ def __init__(self, *conditions: TerminationCondition) -> None: def terminated(self) -> bool: return any(condition.terminated for condition in self._conditions) - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: if self.terminated: raise RuntimeError("Termination condition has already been reached") stop_messages = await asyncio.gather(*[condition(messages) for condition in self._conditions]) - if any(stop_message is not None for stop_message in stop_messages): - content = ", ".join(stop_message.content for stop_message in stop_messages if stop_message is not None) - source = ", ".join(stop_message.source for stop_message in stop_messages if stop_message is not None) + stop_messages_filter = [stop_message for stop_message in stop_messages if stop_message is not None] + if len(stop_messages_filter) > 0: + content = ", ".join(stop_message.content for stop_message in stop_messages_filter) + source = ", ".join(stop_message.source for stop_message in stop_messages_filter) return StopMessage(content=content, source=source) return None diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/__init__.py b/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/__init__.py index 0a53f739848c..72b61745acf1 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/__init__.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/__init__.py @@ -5,6 +5,7 @@ from ._terminations import ( ExternalTermination, + FunctionalTermination, FunctionCallTermination, HandoffTermination, MaxMessageTermination, @@ -27,4 +28,5 @@ "SourceMatchTermination", "TextMessageTermination", "FunctionCallTermination", + "FunctionalTermination", ] diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/_terminations.py b/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/_terminations.py index 7ccddd1f6da4..f0ba274ebe72 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/_terminations.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/_terminations.py @@ -1,5 +1,6 @@ +import asyncio import time -from typing import List, Sequence +from typing import Awaitable, Callable, List, Sequence from autogen_core import Component from pydantic import BaseModel @@ -7,11 +8,9 @@ from ..base import TerminatedException, TerminationCondition from ..messages import ( - AgentEvent, + BaseAgentEvent, BaseChatMessage, - ChatMessage, HandoffMessage, - MultiModalMessage, StopMessage, TextMessage, ToolCallExecutionEvent, @@ -35,7 +34,7 @@ def __init__(self) -> None: def terminated(self) -> bool: return self._terminated - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: if self._terminated: raise TerminatedException("Termination condition has already been reached") for message in messages: @@ -65,8 +64,8 @@ class MaxMessageTermination(TerminationCondition, Component[MaxMessageTerminatio Args: max_messages: The maximum number of messages allowed in the conversation. - include_agent_event: If True, include :class:`~autogen_agentchat.messages.AgentEvent` in the message count. - Otherwise, only include :class:`~autogen_agentchat.messages.ChatMessage`. Defaults to False. + include_agent_event: If True, include :class:`~autogen_agentchat.messages.BaseAgentEvent` in the message count. + Otherwise, only include :class:`~autogen_agentchat.messages.BaseChatMessage`. Defaults to False. """ component_config_schema = MaxMessageTerminationConfig @@ -81,7 +80,7 @@ def __init__(self, max_messages: int, include_agent_event: bool = False) -> None def terminated(self) -> bool: return self._message_count >= self._max_messages - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: if self.terminated: raise TerminatedException("Termination condition has already been reached") self._message_count += len([m for m in messages if self._include_agent_event or isinstance(m, BaseChatMessage)]) @@ -130,25 +129,19 @@ def __init__(self, text: str, sources: Sequence[str] | None = None) -> None: def terminated(self) -> bool: return self._terminated - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: if self._terminated: raise TerminatedException("Termination condition has already been reached") for message in messages: if self._sources is not None and message.source not in self._sources: continue - if isinstance(message.content, str) and self._termination_text in message.content: + content = message.to_text() + if self._termination_text in content: self._terminated = True return StopMessage( content=f"Text '{self._termination_text}' mentioned", source="TextMentionTermination" ) - elif isinstance(message, MultiModalMessage): - for item in message.content: - if isinstance(item, str) and self._termination_text in item: - self._terminated = True - return StopMessage( - content=f"Text '{self._termination_text}' mentioned", source="TextMentionTermination" - ) return None async def reset(self) -> None: @@ -162,6 +155,77 @@ def _from_config(cls, config: TextMentionTerminationConfig) -> Self: return cls(text=config.text) +class FunctionalTermination(TerminationCondition): + """Terminate the conversation if an functional expression is met. + + Args: + func (Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], bool] | Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[bool]]): A function that takes a sequence of messages + and returns True if the termination condition is met, False otherwise. + The function can be a callable or an async callable. + + Example: + + .. code-block:: python + + import asyncio + from typing import Sequence + + from autogen_agentchat.conditions import FunctionalTermination + from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage, StopMessage + + + def expression(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> bool: + # Check if the last message is a stop message + return isinstance(messages[-1], StopMessage) + + + termination = FunctionalTermination(expression) + + + async def run() -> None: + messages = [ + StopMessage(source="agent1", content="Stop"), + ] + result = await termination(messages) + print(result) + + + asyncio.run(run()) + + .. code-block:: text + + StopMessage(source="FunctionalTermination", content="Functional termination condition met") + + """ + + def __init__( + self, + func: Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], bool] + | Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[bool]], + ) -> None: + self._func = func + self._terminated = False + + @property + def terminated(self) -> bool: + return self._terminated + + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: + if self._terminated: + raise TerminatedException("Termination condition has already been reached") + if asyncio.iscoroutinefunction(self._func): + result = await self._func(messages) + else: + result = self._func(messages) + if result is True: + self._terminated = True + return StopMessage(content="Functional termination condition met", source="FunctionalTermination") + return None + + async def reset(self) -> None: + self._terminated = False + + class TokenUsageTerminationConfig(BaseModel): max_total_token: int | None max_prompt_token: int | None @@ -208,7 +272,7 @@ def terminated(self) -> bool: or (self._max_completion_token is not None and self._completion_token_count >= self._max_completion_token) ) - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: if self.terminated: raise TerminatedException("Termination condition has already been reached") for message in messages: @@ -265,7 +329,7 @@ def __init__(self, target: str) -> None: def terminated(self) -> bool: return self._terminated - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: if self._terminated: raise TerminatedException("Termination condition has already been reached") for message in messages: @@ -310,7 +374,7 @@ def __init__(self, timeout_seconds: float) -> None: def terminated(self) -> bool: return self._terminated - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: if self._terminated: raise TerminatedException("Termination condition has already been reached") @@ -372,7 +436,7 @@ def set(self) -> None: """Set the termination condition to terminated.""" self._setted = True - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: if self._terminated: raise TerminatedException("Termination condition has already been reached") if self._setted: @@ -417,7 +481,7 @@ def __init__(self, sources: List[str]) -> None: def terminated(self) -> bool: return self._terminated - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: if self._terminated: raise TerminatedException("Termination condition has already been reached") if not messages: @@ -470,7 +534,7 @@ def __init__(self, source: str | None = None) -> None: def terminated(self) -> bool: return self._terminated - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: if self._terminated: raise TerminatedException("Termination condition has already been reached") for message in messages: @@ -510,6 +574,7 @@ class FunctionCallTermination(TerminationCondition, Component[FunctionCallTermin """ component_config_schema = FunctionCallTerminationConfig + component_provider_override = "autogen_agentchat.conditions.FunctionCallTermination" """The schema for the component configuration.""" def __init__(self, function_name: str) -> None: @@ -520,7 +585,7 @@ def __init__(self, function_name: str) -> None: def terminated(self) -> bool: return self._terminated - async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None: + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: if self._terminated: raise TerminatedException("Termination condition has already been reached") for message in messages: diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py b/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py index 89500a50e344..683a80aa5468 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py @@ -4,18 +4,84 @@ class and includes specific fields relevant to the type of message being sent. """ -from abc import ABC -from typing import Dict, List, Literal +import uuid +from abc import ABC, abstractmethod +from datetime import datetime, timezone +from typing import Any, Dict, Generic, List, Literal, Mapping, Optional, Type, TypeVar -from autogen_core import FunctionCall, Image +from autogen_core import Component, ComponentBase, FunctionCall, Image +from autogen_core.code_executor import CodeBlock, CodeResult from autogen_core.memory import MemoryContent -from autogen_core.models import FunctionExecutionResult, LLMMessage, RequestUsage -from pydantic import BaseModel, ConfigDict, Field -from typing_extensions import Annotated +from autogen_core.models import ( + FunctionExecutionResult, + LLMMessage, + RequestUsage, + UserMessage, +) +from autogen_core.utils import schema_to_pydantic_model +from pydantic import BaseModel, Field, computed_field +from typing_extensions import Annotated, Self class BaseMessage(BaseModel, ABC): - """Base class for all message types.""" + """Abstract base class for all message types in AgentChat. + + .. warning:: + + If you want to create a new message type, do not inherit from this class. + Instead, inherit from :class:`BaseChatMessage` or :class:`BaseAgentEvent` + to clarify the purpose of the message type. + + """ + + @abstractmethod + def to_text(self) -> str: + """Convert the message content to a string-only representation + that can be rendered in the console and inspected by the user or conditions. + This is not used for creating text-only content for models. + For :class:`BaseChatMessage` types, use :meth:`to_model_text` instead.""" + ... + + def dump(self) -> Mapping[str, Any]: + """Convert the message to a JSON-serializable dictionary. + + The default implementation uses the Pydantic model's + :meth:`model_dump` method to convert the message to a dictionary. + Datetime objects are automatically converted to ISO format strings + to ensure JSON serialization compatibility. + Override this method if you want to customize the serialization + process or add additional fields to the output. + """ + return self.model_dump(mode="json") + + @classmethod + def load(cls, data: Mapping[str, Any]) -> Self: + """Create a message from a dictionary of JSON-serializable data. + + The default implementation uses the Pydantic model's + :meth:`model_validate` method to create the message from the data. + Override this method if you want to customize the deserialization + process or add additional fields to the input data.""" + return cls.model_validate(data) + + +class BaseChatMessage(BaseMessage, ABC): + """Abstract base class for chat messages. + + .. note:: + + If you want to create a new message type that is used for agent-to-agent + communication, inherit from this class, or simply use + :class:`StructuredMessage` if your content type is a subclass of + Pydantic BaseModel. + + This class is used for messages that are sent between agents in a chat + conversation. Agents are expected to process the content of the + message using models and return a response as another :class:`BaseChatMessage`. + """ + + id: str = Field(default_factory=lambda: str(uuid.uuid4())) + """Unique identifier for this message.""" source: str """The name of the agent that sent this message.""" @@ -26,27 +92,281 @@ class BaseMessage(BaseModel, ABC): metadata: Dict[str, str] = {} """Additional metadata about the message.""" - model_config = ConfigDict(arbitrary_types_allowed=True) + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + """The time when the message was created.""" + @abstractmethod + def to_model_text(self) -> str: + """Convert the content of the message to text-only representation. + This is used for creating text-only content for models. -class BaseChatMessage(BaseMessage, ABC): - """Base class for chat messages.""" - - pass + This is not used for rendering the message in console. For that, use + :meth:`~BaseMessage.to_text`. + The difference between this and :meth:`to_model_message` is that this + is used to construct parts of the a message for the model client, + while :meth:`to_model_message` is used to create a complete message + for the model client. + """ + ... -class BaseAgentEvent(BaseMessage, ABC): - """Base class for agent events.""" + @abstractmethod + def to_model_message(self) -> UserMessage: + """Convert the message content to a :class:`~autogen_core.models.UserMessage` + for use with model client, e.g., :class:`~autogen_core.models.ChatCompletionClient`. + """ + ... - pass +class BaseTextChatMessage(BaseChatMessage, ABC): + """Base class for all text-only :class:`BaseChatMessage` types. + It has implementations for :meth:`to_text`, :meth:`to_model_text`, + and :meth:`to_model_message` methods. -class TextMessage(BaseChatMessage): - """A text message.""" + Inherit from this class if your message content type is a string. + """ content: str """The content of the message.""" + def to_text(self) -> str: + return self.content + + def to_model_text(self) -> str: + return self.content + + def to_model_message(self) -> UserMessage: + return UserMessage(content=self.content, source=self.source) + + +class BaseAgentEvent(BaseMessage, ABC): + """Base class for agent events. + + .. note:: + + If you want to create a new message type for signaling observable events + to user and application, inherit from this class. + + Agent events are used to signal actions and thoughts produced by agents + and teams to user and applications. They are not used for agent-to-agent + communication and are not expected to be processed by other agents. + + You should override the :meth:`to_text` method if you want to provide + a custom rendering of the content. + """ + + id: str = Field(default_factory=lambda: str(uuid.uuid4())) + """Unique identifier for this event.""" + + source: str + """The name of the agent that sent this message.""" + + models_usage: RequestUsage | None = None + """The model client usage incurred when producing this message.""" + + metadata: Dict[str, str] = {} + """Additional metadata about the message.""" + + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + """The time when the message was created.""" + + +StructuredContentType = TypeVar("StructuredContentType", bound=BaseModel, covariant=True) +"""Type variable for structured content types.""" + + +class StructuredMessage(BaseChatMessage, Generic[StructuredContentType]): + """A :class:`BaseChatMessage` type with an unspecified content type. + + To create a new structured message type, specify the content type + as a subclass of `Pydantic BaseModel `_. + + .. code-block:: python + + from pydantic import BaseModel + from autogen_agentchat.messages import StructuredMessage + + + class MyMessageContent(BaseModel): + text: str + number: int + + + message = StructuredMessage[MyMessageContent]( + content=MyMessageContent(text="Hello", number=42), + source="agent1", + ) + + print(message.to_text()) # {"text": "Hello", "number": 42} + + .. code-block:: python + + from pydantic import BaseModel + from autogen_agentchat.messages import StructuredMessage + + + class MyMessageContent(BaseModel): + text: str + number: int + + + message = StructuredMessage[MyMessageContent]( + content=MyMessageContent(text="Hello", number=42), + source="agent", + format_string="Hello, {text} {number}!", + ) + + print(message.to_text()) # Hello, agent 42! + + """ + + content: StructuredContentType + """The content of the message. Must be a subclass of + `Pydantic BaseModel `_.""" + + format_string: Optional[str] = None + """(Experimental) An optional format string to render the content into a human-readable format. + The format string can use the fields of the content model as placeholders. + For example, if the content model has a field `name`, you can use + `{name}` in the format string to include the value of that field. + The format string is used in the :meth:`to_text` method to create a + human-readable representation of the message. + This setting is experimental and will change in the future. + """ + + @computed_field + def type(self) -> str: + return self.__class__.__name__ + + def to_text(self) -> str: + if self.format_string is not None: + return self.format_string.format(**self.content.model_dump()) + else: + return self.content.model_dump_json() + + def to_model_text(self) -> str: + if self.format_string is not None: + return self.format_string.format(**self.content.model_dump()) + else: + return self.content.model_dump_json() + + def to_model_message(self) -> UserMessage: + return UserMessage( + content=self.content.model_dump_json(), + source=self.source, + ) + + +class StructureMessageConfig(BaseModel): + """The declarative configuration for the structured output.""" + + json_schema: Dict[str, Any] + format_string: Optional[str] = None + content_model_name: str + + +class StructuredMessageFactory(ComponentBase[StructureMessageConfig], Component[StructureMessageConfig]): + """:meta private: + + A component that creates structured chat messages from Pydantic models or JSON schemas. + + This component helps you generate strongly-typed chat messages with content defined using a Pydantic model. + It can be used in declarative workflows where message structure must be validated, formatted, and serialized. + + You can initialize the component directly using a `BaseModel` subclass, or dynamically from a configuration + object (e.g., loaded from disk or a database). + + ### Example 1: Create from a Pydantic Model + + .. code-block:: python + + from pydantic import BaseModel + from autogen_agentchat.messages import StructuredMessageFactory + + + class TestContent(BaseModel): + field1: str + field2: int + + + format_string = "This is a string {field1} and this is an int {field2}" + sm_component = StructuredMessageFactory(input_model=TestContent, format_string=format_string) + + message = sm_component.StructuredMessage( + source="test_agent", content=TestContent(field1="Hello", field2=42), format_string=format_string + ) + + print(message.to_model_text()) # Output: This is a string Hello and this is an int 42 + + config = sm_component.dump_component() + + s_m_dyn = StructuredMessageFactory.load_component(config) + message = s_m_dyn.StructuredMessage( + source="test_agent", + content=s_m_dyn.ContentModel(field1="dyn agent", field2=43), + format_string=s_m_dyn.format_string, + ) + print(type(message)) # StructuredMessage[GeneratedModel] + print(message.to_model_text()) # Output: This is a string dyn agent and this is an int 43 + + Attributes: + component_config_schema (StructureMessageConfig): Defines the configuration structure for this component. + component_provider_override (str): Path used to reference this component in external tooling. + component_type (str): Identifier used for categorization (e.g., "structured_message"). + + Raises: + ValueError: If neither `json_schema` nor `input_model` is provided. + + Args: + json_schema (Optional[str]): JSON schema to dynamically create a Pydantic model. + input_model (Optional[Type[BaseModel]]): A subclass of `BaseModel` that defines the expected message structure. + format_string (Optional[str]): Optional string to render content into a human-readable format. + content_model_name (Optional[str]): Optional name for the generated Pydantic model. + """ + + component_config_schema = StructureMessageConfig + component_provider_override = "autogen_agentchat.messages.StructuredMessageFactory" + component_type = "structured_message" + + def __init__( + self, + json_schema: Optional[Dict[str, Any]] = None, + input_model: Optional[Type[BaseModel]] = None, + format_string: Optional[str] = None, + content_model_name: Optional[str] = None, + ) -> None: + self.format_string = format_string + + if json_schema: + self.ContentModel = schema_to_pydantic_model( + json_schema, model_name=content_model_name or "GeneratedContentModel" + ) + elif input_model: + self.ContentModel = input_model + else: + raise ValueError("Either `json_schema` or `input_model` must be provided.") + + self.StructuredMessage = StructuredMessage[self.ContentModel] # type: ignore[name-defined] + + def _to_config(self) -> StructureMessageConfig: + return StructureMessageConfig( + json_schema=self.ContentModel.model_json_schema(), + format_string=self.format_string, + content_model_name=self.ContentModel.__name__, + ) + + @classmethod + def _from_config(cls, config: StructureMessageConfig) -> "StructuredMessageFactory": + return cls( + json_schema=config.json_schema, + format_string=config.format_string, + content_model_name=config.content_model_name, + ) + + +class TextMessage(BaseTextChatMessage): + """A text message with string-only content.""" + type: Literal["TextMessage"] = "TextMessage" @@ -58,31 +378,70 @@ class MultiModalMessage(BaseChatMessage): type: Literal["MultiModalMessage"] = "MultiModalMessage" - -class StopMessage(BaseChatMessage): + def to_model_text(self, image_placeholder: str | None = "[image]") -> str: + """Convert the content of the message to a string-only representation. + If an image is present, it will be replaced with the image placeholder + by default, otherwise it will be a base64 string when set to None. + """ + text = "" + for c in self.content: + if isinstance(c, str): + text += c + elif isinstance(c, Image): + if image_placeholder is not None: + text += f" {image_placeholder}" + else: + text += f" {c.to_base64()}" + return text + + def to_text(self, iterm: bool = False) -> str: + result: List[str] = [] + for c in self.content: + if isinstance(c, str): + result.append(c) + else: + if iterm: + # iTerm2 image rendering protocol: https://iterm2.com/documentation-images.html + image_data = c.to_base64() + result.append(f"\033]1337;File=inline=1:{image_data}\a\n") + else: + result.append("") + return "\n".join(result) + + def to_model_message(self) -> UserMessage: + return UserMessage(content=self.content, source=self.source) + + +class StopMessage(BaseTextChatMessage): """A message requesting stop of a conversation.""" - content: str - """The content for the stop message.""" - type: Literal["StopMessage"] = "StopMessage" -class HandoffMessage(BaseChatMessage): +class HandoffMessage(BaseTextChatMessage): """A message requesting handoff of a conversation to another agent.""" target: str """The name of the target agent to handoff to.""" - content: str - """The handoff message to the target agent.""" - context: List[LLMMessage] = [] """The model context to be passed to the target agent.""" type: Literal["HandoffMessage"] = "HandoffMessage" +class ToolCallSummaryMessage(BaseTextChatMessage): + """A message signaling the summary of tool call results.""" + + type: Literal["ToolCallSummaryMessage"] = "ToolCallSummaryMessage" + + tool_calls: List[FunctionCall] + """The tool calls that were made.""" + + results: List[FunctionExecutionResult] + """The results of the tool calls.""" + + class ToolCallRequestEvent(BaseAgentEvent): """An event signaling a request to use tools.""" @@ -91,6 +450,42 @@ class ToolCallRequestEvent(BaseAgentEvent): type: Literal["ToolCallRequestEvent"] = "ToolCallRequestEvent" + def to_text(self) -> str: + return str(self.content) + + +class CodeGenerationEvent(BaseAgentEvent): + """An event signaling code generation event.""" + + retry_attempt: int + "Retry number, 0 means first generation" + + content: str + "The complete content as string." + + code_blocks: List[CodeBlock] + "List of code blocks present in content" + + type: Literal["CodeGenerationEvent"] = "CodeGenerationEvent" + + def to_text(self) -> str: + return self.content + + +class CodeExecutionEvent(BaseAgentEvent): + """An event signaling code execution event.""" + + retry_attempt: int + "Retry number, 0 means first execution" + + result: CodeResult + "Code Execution Result" + + type: Literal["CodeExecutionEvent"] = "CodeExecutionEvent" + + def to_text(self) -> str: + return self.result.output + class ToolCallExecutionEvent(BaseAgentEvent): """An event signaling the execution of tool calls.""" @@ -100,14 +495,8 @@ class ToolCallExecutionEvent(BaseAgentEvent): type: Literal["ToolCallExecutionEvent"] = "ToolCallExecutionEvent" - -class ToolCallSummaryMessage(BaseChatMessage): - """A message signaling the summary of tool call results.""" - - content: str - """Summary of the the tool call results.""" - - type: Literal["ToolCallSummaryMessage"] = "ToolCallSummaryMessage" + def to_text(self) -> str: + return str(self.content) class UserInputRequestedEvent(BaseAgentEvent): @@ -121,6 +510,9 @@ class UserInputRequestedEvent(BaseAgentEvent): type: Literal["UserInputRequestedEvent"] = "UserInputRequestedEvent" + def to_text(self) -> str: + return str(self.content) + class MemoryQueryEvent(BaseAgentEvent): """An event signaling the results of memory queries.""" @@ -130,32 +522,134 @@ class MemoryQueryEvent(BaseAgentEvent): type: Literal["MemoryQueryEvent"] = "MemoryQueryEvent" + def to_text(self) -> str: + return str(self.content) + class ModelClientStreamingChunkEvent(BaseAgentEvent): """An event signaling a text output chunk from a model client in streaming mode.""" content: str - """The partial text chunk.""" + """A string chunk from the model client.""" + + full_message_id: str | None = None + """Optional reference to the complete message that may come after the chunks. + This allows consumers of the stream to correlate chunks with the eventual completed message.""" type: Literal["ModelClientStreamingChunkEvent"] = "ModelClientStreamingChunkEvent" + def to_text(self) -> str: + return self.content + class ThoughtEvent(BaseAgentEvent): - """An event signaling the thought process of an agent. + """An event signaling the thought process of a model. It is used to communicate the reasoning tokens generated by a reasoning model, or the extra text content generated by a function call.""" content: str - """The thought process.""" + """The thought process of the model.""" type: Literal["ThoughtEvent"] = "ThoughtEvent" + def to_text(self) -> str: + return self.content + + +class SelectSpeakerEvent(BaseAgentEvent): + """An event signaling the selection of speakers for a conversation.""" + + content: List[str] + """The names of the selected speakers.""" + + type: Literal["SelectSpeakerEvent"] = "SelectSpeakerEvent" + + def to_text(self) -> str: + return str(self.content) + + +class SelectorEvent(BaseAgentEvent): + """An event emitted from the `SelectorGroupChat`.""" + + content: str + """The content of the event.""" + + type: Literal["SelectorEvent"] = "SelectorEvent" + + def to_text(self) -> str: + return str(self.content) + + +class MessageFactory: + """:meta private: + + A factory for creating messages from JSON-serializable dictionaries. + + This is useful for deserializing messages from JSON data. + """ + + def __init__(self) -> None: + self._message_types: Dict[str, type[BaseAgentEvent | BaseChatMessage]] = {} + # Register all message types. + self._message_types[TextMessage.__name__] = TextMessage + self._message_types[MultiModalMessage.__name__] = MultiModalMessage + self._message_types[StopMessage.__name__] = StopMessage + self._message_types[ToolCallSummaryMessage.__name__] = ToolCallSummaryMessage + self._message_types[HandoffMessage.__name__] = HandoffMessage + self._message_types[ToolCallRequestEvent.__name__] = ToolCallRequestEvent + self._message_types[ToolCallExecutionEvent.__name__] = ToolCallExecutionEvent + self._message_types[MemoryQueryEvent.__name__] = MemoryQueryEvent + self._message_types[UserInputRequestedEvent.__name__] = UserInputRequestedEvent + self._message_types[ModelClientStreamingChunkEvent.__name__] = ModelClientStreamingChunkEvent + self._message_types[ThoughtEvent.__name__] = ThoughtEvent + self._message_types[SelectSpeakerEvent.__name__] = SelectSpeakerEvent + self._message_types[CodeGenerationEvent.__name__] = CodeGenerationEvent + self._message_types[CodeExecutionEvent.__name__] = CodeExecutionEvent + + def is_registered(self, message_type: type[BaseAgentEvent | BaseChatMessage]) -> bool: + """Check if a message type is registered with the factory.""" + # Get the class name of the message type. + class_name = message_type.__name__ + # Check if the class name is already registered. + return class_name in self._message_types + + def register(self, message_type: type[BaseAgentEvent | BaseChatMessage]) -> None: + """Register a new message type with the factory.""" + if self.is_registered(message_type): + raise ValueError(f"Message type {message_type} is already registered.") + if not issubclass(message_type, BaseChatMessage) and not issubclass(message_type, BaseAgentEvent): + raise ValueError(f"Message type {message_type} must be a subclass of BaseChatMessage or BaseAgentEvent.") + # Get the class name of the + class_name = message_type.__name__ + # Check if the class name is already registered. + # Register the message type. + self._message_types[class_name] = message_type + + def create(self, data: Mapping[str, Any]) -> BaseAgentEvent | BaseChatMessage: + """Create a message from a dictionary of JSON-serializable data.""" + # Get the type of the message from the dictionary. + message_type = data.get("type") + if message_type is None: + raise ValueError("Field 'type' is required in the message data to recover the message type.") + if message_type not in self._message_types: + raise ValueError(f"Unknown message type: {message_type}") + if not isinstance(message_type, str): + raise ValueError(f"Message type must be a string, got {type(message_type)}") + + # Get the class for the message type. + message_class = self._message_types[message_type] + + # Create an instance of the message class. + assert issubclass(message_class, BaseChatMessage) or issubclass(message_class, BaseAgentEvent) + return message_class.load(data) + ChatMessage = Annotated[ - TextMessage | MultiModalMessage | StopMessage | ToolCallSummaryMessage | HandoffMessage, Field(discriminator="type") + TextMessage | MultiModalMessage | StopMessage | ToolCallSummaryMessage | HandoffMessage, + Field(discriminator="type"), ] -"""Messages for agent-to-agent communication only.""" - +"""The union type of all built-in concrete subclasses of :class:`BaseChatMessage`. +It does not include :class:`StructuredMessage` types.""" AgentEvent = Annotated[ ToolCallRequestEvent @@ -163,16 +657,24 @@ class ThoughtEvent(BaseAgentEvent): | MemoryQueryEvent | UserInputRequestedEvent | ModelClientStreamingChunkEvent - | ThoughtEvent, + | ThoughtEvent + | SelectSpeakerEvent + | CodeGenerationEvent + | CodeExecutionEvent, Field(discriminator="type"), ] -"""Events emitted by agents and teams when they work, not used for agent-to-agent communication.""" - +"""The union type of all built-in concrete subclasses of :class:`BaseAgentEvent`.""" __all__ = [ "AgentEvent", "BaseMessage", "ChatMessage", + "BaseChatMessage", + "BaseAgentEvent", + "BaseTextChatMessage", + "StructuredContentType", + "StructuredMessage", + "StructuredMessageFactory", "HandoffMessage", "MultiModalMessage", "StopMessage", @@ -184,4 +686,8 @@ class ThoughtEvent(BaseAgentEvent): "UserInputRequestedEvent", "ModelClientStreamingChunkEvent", "ThoughtEvent", + "SelectSpeakerEvent", + "MessageFactory", + "CodeGenerationEvent", + "CodeExecutionEvent", ] diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/state/_states.py b/python/packages/autogen-agentchat/src/autogen_agentchat/state/_states.py index 16ddbc7472d6..ecc7b5f7cae7 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/state/_states.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/state/_states.py @@ -1,15 +1,7 @@ -from typing import Annotated, Any, List, Mapping, Optional +from typing import Any, List, Mapping, Optional from pydantic import BaseModel, Field -from ..messages import ( - AgentEvent, - ChatMessage, -) - -# Ensures pydantic can distinguish between types of events & messages. -_AgentMessage = Annotated[AgentEvent | ChatMessage, Field(discriminator="type")] - class BaseState(BaseModel): """Base class for all saveable state""" @@ -35,7 +27,7 @@ class TeamState(BaseState): class BaseGroupChatManagerState(BaseState): """Base state for all group chat managers.""" - message_thread: List[_AgentMessage] = Field(default_factory=list) + message_thread: List[Mapping[str, Any]] = Field(default_factory=list) current_turn: int = Field(default=0) type: str = Field(default="BaseGroupChatManagerState") @@ -44,7 +36,7 @@ class ChatAgentContainerState(BaseState): """State for a container of chat agents.""" agent_state: Mapping[str, Any] = Field(default_factory=dict) - message_buffer: List[ChatMessage] = Field(default_factory=list) + message_buffer: List[Mapping[str, Any]] = Field(default_factory=list) type: str = Field(default="ChatAgentContainerState") diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/__init__.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/__init__.py index e44628edc642..712976980ae5 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/__init__.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/__init__.py @@ -4,6 +4,13 @@ """ from ._group_chat._base_group_chat import BaseGroupChat +from ._group_chat._graph import ( + DiGraph, + DiGraphBuilder, + DiGraphEdge, + DiGraphNode, + GraphFlow, +) from ._group_chat._magentic_one import MagenticOneGroupChat from ._group_chat._round_robin_group_chat import RoundRobinGroupChat from ._group_chat._selector_group_chat import SelectorGroupChat @@ -15,4 +22,9 @@ "SelectorGroupChat", "Swarm", "MagenticOneGroupChat", + "DiGraphBuilder", + "DiGraph", + "DiGraphNode", + "DiGraphEdge", + "GraphFlow", ] diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py index ff6731b03f4a..60f222912387 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py @@ -1,5 +1,4 @@ import asyncio -import logging import uuid from abc import ABC, abstractmethod from typing import Any, AsyncGenerator, Callable, Dict, List, Mapping, Sequence @@ -15,42 +14,70 @@ ) from pydantic import BaseModel, ValidationError -from ... import EVENT_LOGGER_NAME from ...base import ChatAgent, TaskResult, Team, TerminationCondition from ...messages import ( - AgentEvent, + BaseAgentEvent, BaseChatMessage, - ChatMessage, + MessageFactory, ModelClientStreamingChunkEvent, StopMessage, + StructuredMessage, TextMessage, ) from ...state import TeamState from ._chat_agent_container import ChatAgentContainer -from ._events import GroupChatPause, GroupChatReset, GroupChatResume, GroupChatStart, GroupChatTermination +from ._events import ( + GroupChatPause, + GroupChatReset, + GroupChatResume, + GroupChatStart, + GroupChatTermination, + SerializableException, +) from ._sequential_routed_agent import SequentialRoutedAgent -event_logger = logging.getLogger(EVENT_LOGGER_NAME) - class BaseGroupChat(Team, ABC, ComponentBase[BaseModel]): """The base class for group chat teams. + In a group chat team, participants share context by publishing their messages + to all other participants. + + If an :class:`~autogen_agentchat.base.ChatAgent` is a participant, + the :class:`~autogen_agentchat.messages.BaseChatMessage` from the agent response's + :attr:`~autogen_agentchat.base.Response.chat_message` will be published + to other participants in the group chat. + + If a :class:`~autogen_agentchat.base.Team` is a participant, + the :class:`~autogen_agentchat.messages.BaseChatMessage` + from the team result' :attr:`~autogen_agentchat.base.TaskResult.messages` will be published + to other participants in the group chat. + To implement a group chat team, first create a subclass of :class:`BaseGroupChatManager` and then create a subclass of :class:`BaseGroupChat` that uses the group chat manager. + + This base class provides the mapping between the agents of the AgentChat API + and the agent runtime of the Core API, and handles high-level features like + running, pausing, resuming, and resetting the team. """ component_type = "team" def __init__( self, - participants: List[ChatAgent], + name: str, + description: str, + participants: List[ChatAgent | Team], group_chat_manager_name: str, group_chat_manager_class: type[SequentialRoutedAgent], termination_condition: TerminationCondition | None = None, max_turns: int | None = None, runtime: AgentRuntime | None = None, + custom_message_types: List[type[BaseAgentEvent | BaseChatMessage]] | None = None, + emit_team_events: bool = False, ): + self._name = name + self._description = description if len(participants) == 0: raise ValueError("At least one participant is required.") if len(participants) != len(set(participant.name for participant in participants)): @@ -59,6 +86,21 @@ def __init__( self._base_group_chat_manager_class = group_chat_manager_class self._termination_condition = termination_condition self._max_turns = max_turns + self._message_factory = MessageFactory() + if custom_message_types is not None: + for message_type in custom_message_types: + self._message_factory.register(message_type) + + for agent in participants: + if isinstance(agent, ChatAgent): + for message_type in agent.produced_message_types: + try: + is_registered = self._message_factory.is_registered(message_type) # type: ignore[reportUnknownArgumentType] + if issubclass(message_type, StructuredMessage) and not is_registered: + self._message_factory.register(message_type) # type: ignore[reportUnknownArgumentType] + except TypeError: + # Not a class or not a valid subclassable type (skip) + pass # The team ID is a UUID that is used to identify the team and its participants # in the agent runtime. It is used to create unique topic types for each participant. @@ -85,7 +127,9 @@ def __init__( self._output_topic_type = f"output_topic_{self._team_id}" # The queue for collecting the output messages. - self._output_message_queue: asyncio.Queue[AgentEvent | ChatMessage | GroupChatTermination] = asyncio.Queue() + self._output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination] = ( + asyncio.Queue() + ) # Create a runtime for the team. if runtime is not None: @@ -103,6 +147,19 @@ def __init__( # Flag to track if the group chat is running. self._is_running = False + # Flag to track if the team events should be emitted. + self._emit_team_events = emit_team_events + + @property + def name(self) -> str: + """The name of the group chat team.""" + return self._name + + @property + def description(self) -> str: + """A description of the group chat team.""" + return self._description + @abstractmethod def _create_group_chat_manager_factory( self, @@ -112,19 +169,21 @@ def _create_group_chat_manager_factory( participant_topic_types: List[str], participant_names: List[str], participant_descriptions: List[str], - output_message_queue: asyncio.Queue[AgentEvent | ChatMessage | GroupChatTermination], + output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination], termination_condition: TerminationCondition | None, max_turns: int | None, + message_factory: MessageFactory, ) -> Callable[[], SequentialRoutedAgent]: ... def _create_participant_factory( self, parent_topic_type: str, output_topic_type: str, - agent: ChatAgent, + agent: ChatAgent | Team, + message_factory: MessageFactory, ) -> Callable[[], ChatAgentContainer]: def _factory() -> ChatAgentContainer: - container = ChatAgentContainer(parent_topic_type, output_topic_type, agent) + container = ChatAgentContainer(parent_topic_type, output_topic_type, agent, message_factory) return container return _factory @@ -140,7 +199,9 @@ async def _init(self, runtime: AgentRuntime) -> None: await ChatAgentContainer.register( runtime, type=agent_type, - factory=self._create_participant_factory(self._group_topic_type, self._output_topic_type, participant), + factory=self._create_participant_factory( + self._group_topic_type, self._output_topic_type, participant, self._message_factory + ), ) # Add subscriptions for the participant. # The participant should be able to receive messages from its own topic. @@ -162,6 +223,7 @@ async def _init(self, runtime: AgentRuntime) -> None: output_message_queue=self._output_message_queue, termination_condition=self._termination_condition, max_turns=self._max_turns, + message_factory=self._message_factory, ), ) # Add subscriptions for the group chat manager. @@ -185,15 +247,16 @@ async def _init(self, runtime: AgentRuntime) -> None: async def run( self, *, - task: str | ChatMessage | Sequence[ChatMessage] | None = None, + task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None, cancellation_token: CancellationToken | None = None, + output_task_messages: bool = True, ) -> TaskResult: """Run the team and return the result. The base implementation uses :meth:`run_stream` to run the team and then returns the final result. Once the team is stopped, the termination condition is reset. Args: - task (str | ChatMessage | Sequence[ChatMessage] | None): The task to run the team with. Can be a string, a single :class:`ChatMessage` , or a list of :class:`ChatMessage`. + task (str | BaseChatMessage | Sequence[BaseChatMessage] | None): The task to run the team with. Can be a string, a single :class:`BaseChatMessage` , or a list of :class:`BaseChatMessage`. cancellation_token (CancellationToken | None): The cancellation token to kill the task immediately. Setting the cancellation token potentially put the team in an inconsistent state, and it may not reset the termination condition. @@ -277,6 +340,7 @@ async def main() -> None: async for message in self.run_stream( task=task, cancellation_token=cancellation_token, + output_task_messages=output_task_messages, ): if isinstance(message, TaskResult): result = message @@ -287,9 +351,10 @@ async def main() -> None: async def run_stream( self, *, - task: str | ChatMessage | Sequence[ChatMessage] | None = None, + task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None, cancellation_token: CancellationToken | None = None, - ) -> AsyncGenerator[AgentEvent | ChatMessage | TaskResult, None]: + output_task_messages: bool = True, + ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None]: """Run the team and produces a stream of messages and the final result of the type :class:`~autogen_agentchat.base.TaskResult` as the last item in the stream. Once the team is stopped, the termination condition is reset. @@ -301,14 +366,15 @@ async def run_stream( :attr:`~autogen_agentchat.base.TaskResult.messages`. Args: - task (str | ChatMessage | Sequence[ChatMessage] | None): The task to run the team with. Can be a string, a single :class:`ChatMessage` , or a list of :class:`ChatMessage`. + task (str | BaseChatMessage | Sequence[BaseChatMessage] | None): The task to run the team with. Can be a string, a single :class:`BaseChatMessage` , or a list of :class:`BaseChatMessage`. cancellation_token (CancellationToken | None): The cancellation token to kill the task immediately. Setting the cancellation token potentially put the team in an inconsistent state, and it may not reset the termination condition. To gracefully stop the team, use :class:`~autogen_agentchat.conditions.ExternalTermination` instead. + output_task_messages (bool): Whether to include task messages in the output stream. Defaults to True for backward compatibility. Returns: - stream: an :class:`~collections.abc.AsyncGenerator` that yields :class:`~autogen_agentchat.messages.AgentEvent`, :class:`~autogen_agentchat.messages.ChatMessage`, and the final result :class:`~autogen_agentchat.base.TaskResult` as the last item in the stream. + stream: an :class:`~collections.abc.AsyncGenerator` that yields :class:`~autogen_agentchat.messages.BaseAgentEvent`, :class:`~autogen_agentchat.messages.BaseChatMessage`, and the final result :class:`~autogen_agentchat.base.TaskResult` as the last item in the stream. Example using the :class:`~autogen_agentchat.teams.RoundRobinGroupChat` team: @@ -386,23 +452,33 @@ async def main() -> None: asyncio.run(main()) """ - # Create the messages list if the task is a string or a chat message. - messages: List[ChatMessage] | None = None + messages: List[BaseChatMessage] | None = None if task is None: pass elif isinstance(task, str): messages = [TextMessage(content=task, source="user")] elif isinstance(task, BaseChatMessage): messages = [task] - else: + elif isinstance(task, list): if not task: raise ValueError("Task list cannot be empty.") messages = [] for msg in task: if not isinstance(msg, BaseChatMessage): - raise ValueError("All messages in task list must be valid ChatMessage types") + raise ValueError("All messages in task list must be valid BaseChatMessage types") messages.append(msg) + else: + raise ValueError("Task must be a string, a BaseChatMessage, or a list of BaseChatMessage.") + # Check if the messages types are registered with the message factory. + if messages is not None: + for msg in messages: + if not self._message_factory.is_registered(msg.__class__): + raise ValueError( + f"Message type {msg.__class__} is not registered with the message factory. " + "Please register it with the message factory by adding it to the " + "custom_message_types list when creating the team." + ) if self._is_running: raise ValueError("The team is already running, it cannot run again until it is stopped.") @@ -424,13 +500,26 @@ async def stop_runtime() -> None: try: # This will propagate any exceptions raised. await self._runtime.stop_when_idle() - finally: + # Put a termination message in the queue to indicate that the group chat is stopped for whatever reason + # but not due to an exception. + await self._output_message_queue.put( + GroupChatTermination( + message=StopMessage( + content="The group chat is stopped.", source=self._group_chat_manager_name + ) + ) + ) + except Exception as e: # Stop the consumption of messages and end the stream. - # NOTE: we also need to put a GroupChatTermination event here because when the group chat + # NOTE: we also need to put a GroupChatTermination event here because when the runtime # has an exception, the group chat manager may not be able to put a GroupChatTermination event in the queue. + # This may not be necessary if the group chat manager is able to handle the exception and put the event in the queue. await self._output_message_queue.put( GroupChatTermination( - message=StopMessage(content="Exception occurred.", source=self._group_chat_manager_name) + message=StopMessage( + content="An exception occurred in the runtime.", source=self._group_chat_manager_name + ), + error=SerializableException.from_exception(e), ) ) @@ -443,14 +532,15 @@ async def stop_runtime() -> None: # The group chat manager will start the group chat by relaying the message to the participants # and the group chat manager. await self._runtime.send_message( - GroupChatStart(messages=messages), + GroupChatStart(messages=messages, output_task_messages=output_task_messages), recipient=AgentId(type=self._group_chat_manager_topic_type, key=self._team_id), cancellation_token=cancellation_token, ) # Collect the output messages in order. - output_messages: List[AgentEvent | ChatMessage] = [] + output_messages: List[BaseAgentEvent | BaseChatMessage] = [] stop_reason: str | None = None - # Yield the messsages until the queue is empty. + + # Yield the messages until the queue is empty. while True: message_future = asyncio.ensure_future(self._output_message_queue.get()) if cancellation_token is not None: @@ -458,11 +548,10 @@ async def stop_runtime() -> None: # Wait for the next message, this will raise an exception if the task is cancelled. message = await message_future if isinstance(message, GroupChatTermination): - # If the message is None, it means the group chat has terminated. - # TODO: how do we handle termination when the runtime is not embedded - # and there is an exception in the group chat? - # The group chat manager may not be able to put a GroupChatTermination event in the queue, - # and this loop will never end. + # If the message contains an error, we need to raise it here. + # This will stop the team and propagate the error. + if message.error is not None: + raise RuntimeError(str(message.error)) stop_reason = message.message.content break yield message diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat_manager.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat_manager.py index 0cd45633728a..b0a0c1d55fc4 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat_manager.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat_manager.py @@ -1,20 +1,23 @@ import asyncio from abc import ABC, abstractmethod -from typing import Any, List +from typing import Any, List, Sequence -from autogen_core import DefaultTopicId, MessageContext, event, rpc +from autogen_core import CancellationToken, DefaultTopicId, MessageContext, event, rpc from ...base import TerminationCondition -from ...messages import AgentEvent, ChatMessage, StopMessage +from ...messages import BaseAgentEvent, BaseChatMessage, MessageFactory, SelectSpeakerEvent, StopMessage from ._events import ( GroupChatAgentResponse, + GroupChatError, GroupChatMessage, GroupChatPause, GroupChatRequestPublish, GroupChatReset, GroupChatResume, GroupChatStart, + GroupChatTeamResponse, GroupChatTermination, + SerializableException, ) from ._sequential_routed_agent import SequentialRoutedAgent @@ -39,40 +42,46 @@ def __init__( participant_topic_types: List[str], participant_names: List[str], participant_descriptions: List[str], - output_message_queue: asyncio.Queue[AgentEvent | ChatMessage | GroupChatTermination], - termination_condition: TerminationCondition | None = None, - max_turns: int | None = None, + output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination], + termination_condition: TerminationCondition | None, + max_turns: int | None, + message_factory: MessageFactory, + emit_team_events: bool = False, ): super().__init__( description="Group chat manager", sequential_message_types=[ GroupChatStart, GroupChatAgentResponse, + GroupChatTeamResponse, GroupChatMessage, GroupChatReset, ], ) - self._name = name - self._group_topic_type = group_topic_type - self._output_topic_type = output_topic_type + if max_turns is not None and max_turns <= 0: + raise ValueError("The maximum number of turns must be greater than 0.") if len(participant_topic_types) != len(participant_descriptions): raise ValueError("The number of participant topic types, agent types, and descriptions must be the same.") if len(set(participant_topic_types)) != len(participant_topic_types): raise ValueError("The participant topic types must be unique.") if group_topic_type in participant_topic_types: raise ValueError("The group topic type must not be in the participant topic types.") + self._name = name + self._group_topic_type = group_topic_type + self._output_topic_type = output_topic_type self._participant_names = participant_names self._participant_name_to_topic_type = { name: topic_type for name, topic_type in zip(participant_names, participant_topic_types, strict=True) } self._participant_descriptions = participant_descriptions - self._message_thread: List[AgentEvent | ChatMessage] = [] + self._message_thread: List[BaseAgentEvent | BaseChatMessage] = [] self._output_message_queue = output_message_queue self._termination_condition = termination_condition - if max_turns is not None and max_turns <= 0: - raise ValueError("The maximum number of turns must be greater than 0.") self._max_turns = max_turns self._current_turn = 0 + self._message_factory = message_factory + self._emit_team_events = emit_team_events + self._active_speakers: List[str] = [] @rpc async def handle_start(self, message: GroupChatStart, ctx: MessageContext) -> None: @@ -98,8 +107,11 @@ async def handle_start(self, message: GroupChatStart, ctx: MessageContext) -> No GroupChatStart(messages=message.messages), topic_id=DefaultTopicId(type=self._output_topic_type), ) - for msg in message.messages: - await self._output_message_queue.put(msg) + + # Only put messages in output queue if output_task_messages is True + if message.output_task_messages: + for msg in message.messages: + await self._output_message_queue.put(msg) # Relay all messages at once to participants await self.publish_message( @@ -109,45 +121,82 @@ async def handle_start(self, message: GroupChatStart, ctx: MessageContext) -> No ) # Append all messages to thread - self._message_thread.extend(message.messages) + await self.update_message_thread(message.messages) # Check termination condition after processing all messages - if self._termination_condition is not None: - stop_message = await self._termination_condition(message.messages) - if stop_message is not None: - # Reset the termination condition. - await self._termination_condition.reset() - # Signal termination to the caller of the team. - await self._signal_termination(stop_message) - # Stop the group chat. - return + if await self._apply_termination_condition(message.messages): + # Stop the group chat. + return - # Select a speaker to start/continue the conversation - speaker_name_future = asyncio.ensure_future(self.select_speaker(self._message_thread)) - # Link the select speaker future to the cancellation token. - ctx.cancellation_token.link_future(speaker_name_future) - speaker_name = await speaker_name_future - if speaker_name not in self._participant_name_to_topic_type: - raise RuntimeError(f"Speaker {speaker_name} not found in participant names.") - speaker_topic_type = self._participant_name_to_topic_type[speaker_name] - await self.publish_message( - GroupChatRequestPublish(), - topic_id=DefaultTopicId(type=speaker_topic_type), - cancellation_token=ctx.cancellation_token, - ) + # Select speakers to start/continue the conversation + await self._transition_to_next_speakers(ctx.cancellation_token) @event - async def handle_agent_response(self, message: GroupChatAgentResponse, ctx: MessageContext) -> None: - # Append the message to the message thread and construct the delta. - delta: List[AgentEvent | ChatMessage] = [] - if message.agent_response.inner_messages is not None: - for inner_message in message.agent_response.inner_messages: - self._message_thread.append(inner_message) - delta.append(inner_message) - self._message_thread.append(message.agent_response.chat_message) - delta.append(message.agent_response.chat_message) - - # Check if the conversation should be terminated. + async def handle_agent_response( + self, message: GroupChatAgentResponse | GroupChatTeamResponse, ctx: MessageContext + ) -> None: + try: + # Construct the detla from the agent response. + delta: List[BaseAgentEvent | BaseChatMessage] = [] + if isinstance(message, GroupChatAgentResponse): + if message.response.inner_messages is not None: + for inner_message in message.response.inner_messages: + delta.append(inner_message) + delta.append(message.response.chat_message) + else: + delta.extend(message.result.messages) + + # Append the messages to the message thread. + await self.update_message_thread(delta) + + # Remove the agent from the active speakers list. + self._active_speakers.remove(message.name) + if len(self._active_speakers) > 0: + # If there are still active speakers, return without doing anything. + return + + # Check if the conversation should be terminated. + if await self._apply_termination_condition(delta, increment_turn_count=True): + # Stop the group chat. + return + + # Select speakers to continue the conversation. + await self._transition_to_next_speakers(ctx.cancellation_token) + except Exception as e: + # Handle the exception and signal termination with an error. + error = SerializableException.from_exception(e) + await self._signal_termination_with_error(error) + # Raise the exception to the runtime. + raise + + async def _transition_to_next_speakers(self, cancellation_token: CancellationToken) -> None: + speaker_names_future = asyncio.ensure_future(self.select_speaker(self._message_thread)) + # Link the select speaker future to the cancellation token. + cancellation_token.link_future(speaker_names_future) + speaker_names = await speaker_names_future + if isinstance(speaker_names, str): + # If only one speaker is selected, convert it to a list. + speaker_names = [speaker_names] + for speaker_name in speaker_names: + if speaker_name not in self._participant_name_to_topic_type: + raise RuntimeError(f"Speaker {speaker_name} not found in participant names.") + await self._log_speaker_selection(speaker_names) + + # Send request to publish message to the next speakers + for speaker_name in speaker_names: + speaker_topic_type = self._participant_name_to_topic_type[speaker_name] + await self.publish_message( + GroupChatRequestPublish(), + topic_id=DefaultTopicId(type=speaker_topic_type), + cancellation_token=cancellation_token, + ) + self._active_speakers.append(speaker_name) + + async def _apply_termination_condition( + self, delta: Sequence[BaseAgentEvent | BaseChatMessage], increment_turn_count: bool = False + ) -> bool: + """Apply the termination condition to the delta and return True if the conversation should be terminated. + It also resets the termination condition and turn count, and signals termination to the caller of the team.""" if self._termination_condition is not None: stop_message = await self._termination_condition(delta) if stop_message is not None: @@ -157,10 +206,10 @@ async def handle_agent_response(self, message: GroupChatAgentResponse, ctx: Mess # Signal termination to the caller of the team. await self._signal_termination(stop_message) # Stop the group chat. - return - - # Increment the turn count. - self._current_turn += 1 + return True + if increment_turn_count: + # Increment the turn count. + self._current_turn += 1 # Check if the maximum number of turns has been reached. if self._max_turns is not None: if self._current_turn >= self._max_turns: @@ -175,21 +224,18 @@ async def handle_agent_response(self, message: GroupChatAgentResponse, ctx: Mess # Signal termination to the caller of the team. await self._signal_termination(stop_message) # Stop the group chat. - return + return True + return False - # Select a speaker to continue the conversation. - speaker_name_future = asyncio.ensure_future(self.select_speaker(self._message_thread)) - # Link the select speaker future to the cancellation token. - ctx.cancellation_token.link_future(speaker_name_future) - speaker_name = await speaker_name_future - if speaker_name not in self._participant_name_to_topic_type: - raise RuntimeError(f"Speaker {speaker_name} not found in participant names.") - speaker_topic_type = self._participant_name_to_topic_type[speaker_name] - await self.publish_message( - GroupChatRequestPublish(), - topic_id=DefaultTopicId(type=speaker_topic_type), - cancellation_token=ctx.cancellation_token, - ) + async def _log_speaker_selection(self, speaker_names: List[str]) -> None: + """Log the selected speaker to the output message queue.""" + select_msg = SelectSpeakerEvent(content=speaker_names, source=self._name) + if self._emit_team_events: + await self.publish_message( + GroupChatMessage(message=select_msg), + topic_id=DefaultTopicId(type=self._output_topic_type), + ) + await self._output_message_queue.put(select_msg) async def _signal_termination(self, message: StopMessage) -> None: termination_event = GroupChatTermination(message=message) @@ -201,11 +247,28 @@ async def _signal_termination(self, message: StopMessage) -> None: # Put the termination event in the output message queue. await self._output_message_queue.put(termination_event) + async def _signal_termination_with_error(self, error: SerializableException) -> None: + termination_event = GroupChatTermination( + message=StopMessage(content="An error occurred in the group chat.", source=self._name), error=error + ) + # Log the termination event. + await self.publish_message( + termination_event, + topic_id=DefaultTopicId(type=self._output_topic_type), + ) + # Put the termination event in the output message queue. + await self._output_message_queue.put(termination_event) + @event async def handle_group_chat_message(self, message: GroupChatMessage, ctx: MessageContext) -> None: """Handle a group chat message by appending the content to its output message queue.""" await self._output_message_queue.put(message.message) + @event + async def handle_group_chat_error(self, message: GroupChatError, ctx: MessageContext) -> None: + """Handle a group chat error by logging the error and signaling termination.""" + await self._signal_termination_with_error(message.error) + @rpc async def handle_reset(self, message: GroupChatReset, ctx: MessageContext) -> None: """Reset the group chat manager. Calling :meth:`reset` to reset the group chat manager @@ -223,7 +286,7 @@ async def handle_resume(self, message: GroupChatResume, ctx: MessageContext) -> pass @abstractmethod - async def validate_group_state(self, messages: List[ChatMessage] | None) -> None: + async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None: """Validate the state of the group chat given the start messages. This is executed when the group chat manager receives a GroupChatStart event. @@ -232,10 +295,26 @@ async def validate_group_state(self, messages: List[ChatMessage] | None) -> None """ ... + async def update_message_thread(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> None: + """Update the message thread with the new messages. + This is called when the group chat receives a GroupChatStart or GroupChatAgentResponse event, + before calling the select_speakers method. + """ + self._message_thread.extend(messages) + @abstractmethod - async def select_speaker(self, thread: List[AgentEvent | ChatMessage]) -> str: - """Select a speaker from the participants and return the - topic type of the selected speaker.""" + async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str: + """Select speakers from the participants and return the topic types of the selected speaker. + This is called when the group chat manager have received all responses from the participants + for a turn and is ready to select the next speakers for the next turn. + + Args: + thread: The message thread of the group chat. + + Returns: + A list of topic types of the selected speakers. + If only one speaker is selected, a single string is returned instead of a list. + """ ... @abstractmethod diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_chat_agent_container.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_chat_agent_container.py index 7c86556e257a..ff660c6f78e3 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_chat_agent_container.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_chat_agent_container.py @@ -1,34 +1,42 @@ from typing import Any, List, Mapping -from autogen_core import DefaultTopicId, MessageContext, event, rpc +from autogen_core import DefaultTopicId, MessageContext, event, rpc, trace_invoke_agent_span -from ...base import ChatAgent, Response -from ...messages import ChatMessage +from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage, MessageFactory + +from ...base import ChatAgent, Response, TaskResult, Team from ...state import ChatAgentContainerState from ._events import ( GroupChatAgentResponse, + GroupChatError, GroupChatMessage, GroupChatPause, GroupChatRequestPublish, GroupChatReset, GroupChatResume, GroupChatStart, + GroupChatTeamResponse, + SerializableException, ) from ._sequential_routed_agent import SequentialRoutedAgent class ChatAgentContainer(SequentialRoutedAgent): """A core agent class that delegates message handling to an - :class:`autogen_agentchat.base.ChatAgent` so that it can be used in a - group chat team. + :class:`autogen_agentchat.base.ChatAgent` or :class:`autogen_agentchat.base.Team` + so that it can be used in a group chat team. Args: parent_topic_type (str): The topic type of the parent orchestrator. output_topic_type (str): The topic type for the output. - agent (ChatAgent): The agent to delegate message handling to. + agent (ChatAgent | Team): The agent or team to delegate message handling to. + message_factory (MessageFactory): The message factory to use for + creating messages from JSON data. """ - def __init__(self, parent_topic_type: str, output_topic_type: str, agent: ChatAgent) -> None: + def __init__( + self, parent_topic_type: str, output_topic_type: str, agent: ChatAgent | Team, message_factory: MessageFactory + ) -> None: super().__init__( description=agent.description, sequential_message_types=[ @@ -36,79 +44,170 @@ def __init__(self, parent_topic_type: str, output_topic_type: str, agent: ChatAg GroupChatRequestPublish, GroupChatReset, GroupChatAgentResponse, + GroupChatTeamResponse, ], ) self._parent_topic_type = parent_topic_type self._output_topic_type = output_topic_type self._agent = agent - self._message_buffer: List[ChatMessage] = [] + self._message_buffer: List[BaseChatMessage] = [] + self._message_factory = message_factory @event async def handle_start(self, message: GroupChatStart, ctx: MessageContext) -> None: """Handle a start event by appending the content to the buffer.""" if message.messages is not None: - self._message_buffer.extend(message.messages) + for msg in message.messages: + self._buffer_message(msg) @event async def handle_agent_response(self, message: GroupChatAgentResponse, ctx: MessageContext) -> None: """Handle an agent response event by appending the content to the buffer.""" - self._message_buffer.append(message.agent_response.chat_message) + self._buffer_message(message.response.chat_message) + + @event + async def handle_team_response(self, message: GroupChatTeamResponse, ctx: MessageContext) -> None: + """Handle a team response event by appending the content to the buffer.""" + for msg in message.result.messages: + if isinstance(msg, BaseChatMessage): + self._buffer_message(msg) @rpc async def handle_reset(self, message: GroupChatReset, ctx: MessageContext) -> None: """Handle a reset event by resetting the agent.""" self._message_buffer.clear() - await self._agent.on_reset(ctx.cancellation_token) + if isinstance(self._agent, Team): + # If the agent is a team, reset the team. + await self._agent.reset() + else: + await self._agent.on_reset(ctx.cancellation_token) @event async def handle_request(self, message: GroupChatRequestPublish, ctx: MessageContext) -> None: """Handle a content request event by passing the messages in the buffer to the delegate agent and publish the response.""" - # Pass the messages in the buffer to the delegate agent. - response: Response | None = None - async for msg in self._agent.on_messages_stream(self._message_buffer, ctx.cancellation_token): - if isinstance(msg, Response): - # Log the response. + if isinstance(self._agent, Team): + try: + stream = self._agent.run_stream( + task=self._message_buffer, + cancellation_token=ctx.cancellation_token, + output_task_messages=False, + ) + result: TaskResult | None = None + async for team_event in stream: + if isinstance(team_event, TaskResult): + result = team_event + else: + await self._log_message(team_event) + if result is None: + raise RuntimeError( + "The team did not produce a final TaskResult. Check the team's run_stream method." + ) + self._message_buffer.clear() + # Publish the team response to the group chat. await self.publish_message( - GroupChatMessage(message=msg.chat_message), - topic_id=DefaultTopicId(type=self._output_topic_type), + GroupChatTeamResponse(result=result, name=self._agent.name), + topic_id=DefaultTopicId(type=self._parent_topic_type), + cancellation_token=ctx.cancellation_token, ) - response = msg - else: - # Log the message. + except Exception as e: + # Publish the error to the group chat. + error_message = SerializableException.from_exception(e) await self.publish_message( - GroupChatMessage(message=msg), topic_id=DefaultTopicId(type=self._output_topic_type) + GroupChatError(error=error_message), + topic_id=DefaultTopicId(type=self._parent_topic_type), + cancellation_token=ctx.cancellation_token, ) - if response is None: - raise ValueError("The agent did not produce a final response. Check the agent's on_messages_stream method.") - - # Publish the response to the group chat. - self._message_buffer.clear() + # Raise the error to the runtime. + raise + else: + # If the agent is not a team, handle it as a single agent. + with trace_invoke_agent_span( + agent_name=self._agent.name, + agent_description=self._agent.description, + agent_id=str(self.id), + ): + try: + # Pass the messages in the buffer to the delegate agent. + response: Response | None = None + async for msg in self._agent.on_messages_stream(self._message_buffer, ctx.cancellation_token): + if isinstance(msg, Response): + await self._log_message(msg.chat_message) + response = msg + else: + await self._log_message(msg) + if response is None: + raise RuntimeError( + "The agent did not produce a final response. Check the agent's on_messages_stream method." + ) + # Publish the response to the group chat. + self._message_buffer.clear() + await self.publish_message( + GroupChatAgentResponse(response=response, name=self._agent.name), + topic_id=DefaultTopicId(type=self._parent_topic_type), + cancellation_token=ctx.cancellation_token, + ) + except Exception as e: + # Publish the error to the group chat. + error_message = SerializableException.from_exception(e) + await self.publish_message( + GroupChatError(error=error_message), + topic_id=DefaultTopicId(type=self._parent_topic_type), + cancellation_token=ctx.cancellation_token, + ) + # Raise the error to the runtime. + raise + + def _buffer_message(self, message: BaseChatMessage) -> None: + if not self._message_factory.is_registered(message.__class__): + raise ValueError(f"Message type {message.__class__} is not registered.") + # Buffer the message. + self._message_buffer.append(message) + + async def _log_message(self, message: BaseAgentEvent | BaseChatMessage) -> None: + if not self._message_factory.is_registered(message.__class__): + raise ValueError(f"Message type {message.__class__} is not registered.") + # Log the message. await self.publish_message( - GroupChatAgentResponse(agent_response=response), - topic_id=DefaultTopicId(type=self._parent_topic_type), - cancellation_token=ctx.cancellation_token, + GroupChatMessage(message=message), + topic_id=DefaultTopicId(type=self._output_topic_type), ) @rpc async def handle_pause(self, message: GroupChatPause, ctx: MessageContext) -> None: """Handle a pause event by pausing the agent.""" - await self._agent.on_pause(ctx.cancellation_token) + if isinstance(self._agent, Team): + # If the agent is a team, pause the team. + await self._agent.pause() + else: + await self._agent.on_pause(ctx.cancellation_token) @rpc async def handle_resume(self, message: GroupChatResume, ctx: MessageContext) -> None: """Handle a resume event by resuming the agent.""" - await self._agent.on_resume(ctx.cancellation_token) + if isinstance(self._agent, Team): + # If the agent is a team, resume the team. + await self._agent.resume() + else: + await self._agent.on_resume(ctx.cancellation_token) async def on_unhandled_message(self, message: Any, ctx: MessageContext) -> None: raise ValueError(f"Unhandled message in agent container: {type(message)}") async def save_state(self) -> Mapping[str, Any]: agent_state = await self._agent.save_state() - state = ChatAgentContainerState(agent_state=agent_state, message_buffer=list(self._message_buffer)) + state = ChatAgentContainerState( + agent_state=agent_state, message_buffer=[message.dump() for message in self._message_buffer] + ) return state.model_dump() async def load_state(self, state: Mapping[str, Any]) -> None: container_state = ChatAgentContainerState.model_validate(state) - self._message_buffer = list(container_state.message_buffer) + self._message_buffer = [] + for message_data in container_state.message_buffer: + message = self._message_factory.create(message_data) + if isinstance(message, BaseChatMessage): + self._message_buffer.append(message) + else: + raise ValueError(f"Invalid message type in message buffer: {type(message)}") await self._agent.load_state(container_state.agent_state) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_events.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_events.py index f705a54c4a50..a149e5861c27 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_events.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_events.py @@ -1,24 +1,69 @@ +import traceback from typing import List -from pydantic import BaseModel +from pydantic import BaseModel, SerializeAsAny -from ...base import Response -from ...messages import AgentEvent, ChatMessage, StopMessage +from ...base import Response, TaskResult +from ...messages import BaseAgentEvent, BaseChatMessage, StopMessage + + +class SerializableException(BaseModel): + """A serializable exception.""" + + error_type: str + """The type of error that occurred.""" + + error_message: str + """The error message that describes the error.""" + + traceback: str | None = None + """The traceback of the error, if available.""" + + @classmethod + def from_exception(cls, exc: Exception) -> "SerializableException": + """Create a GroupChatError from an exception.""" + return cls( + error_type=type(exc).__name__, + error_message=str(exc), + traceback="\n".join(traceback.format_exception(type(exc), exc, exc.__traceback__)), + ) + + def __str__(self) -> str: + """Return a string representation of the error, including the traceback if available.""" + if self.traceback: + return f"{self.error_type}: {self.error_message}\nTraceback:\n{self.traceback}" + return f"{self.error_type}: {self.error_message}" class GroupChatStart(BaseModel): """A request to start a group chat.""" - messages: List[ChatMessage] | None = None + messages: List[SerializeAsAny[BaseChatMessage]] | None = None """An optional list of messages to start the group chat.""" + output_task_messages: bool = True + """Whether to include task messages in the output. Defaults to True for backward compatibility.""" + class GroupChatAgentResponse(BaseModel): """A response published to a group chat.""" - agent_response: Response + response: SerializeAsAny[Response] """The response from an agent.""" + name: str + """The name of the agent that produced the response.""" + + +class GroupChatTeamResponse(BaseModel): + """A response published to a group chat from a team.""" + + result: SerializeAsAny[TaskResult] + """The result from a team.""" + + name: str + """The name of the team that produced the response.""" + class GroupChatRequestPublish(BaseModel): """A request to publish a message to a group chat.""" @@ -29,7 +74,7 @@ class GroupChatRequestPublish(BaseModel): class GroupChatMessage(BaseModel): """A message from a group chat.""" - message: AgentEvent | ChatMessage + message: SerializeAsAny[BaseAgentEvent | BaseChatMessage] """The message that was published.""" @@ -39,6 +84,9 @@ class GroupChatTermination(BaseModel): message: StopMessage """The stop message that indicates the reason of termination.""" + error: SerializableException | None = None + """The error that occurred, if any.""" + class GroupChatReset(BaseModel): """A request to reset the agents in the group chat.""" @@ -56,3 +104,10 @@ class GroupChatResume(BaseModel): """A request to resume the group chat.""" ... + + +class GroupChatError(BaseModel): + """A message indicating that an error occurred in the group chat.""" + + error: SerializableException + """The error that occurred.""" diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/__init__.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/__init__.py new file mode 100644 index 000000000000..f38d6d653ca8 --- /dev/null +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/__init__.py @@ -0,0 +1,17 @@ +from ._digraph_group_chat import ( + DiGraph, + DiGraphEdge, + DiGraphNode, + GraphFlow, + GraphFlowManager, +) +from ._graph_builder import DiGraphBuilder + +__all__ = [ + "GraphFlow", + "DiGraph", + "GraphFlowManager", + "DiGraphNode", + "DiGraphEdge", + "DiGraphBuilder", +] diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/_digraph_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/_digraph_group_chat.py new file mode 100644 index 000000000000..b9b607b22dd9 --- /dev/null +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/_digraph_group_chat.py @@ -0,0 +1,877 @@ +import asyncio +from collections import Counter, deque +from typing import Any, Callable, Deque, Dict, List, Literal, Mapping, Sequence, Set, Union + +from autogen_core import AgentRuntime, Component, ComponentModel +from pydantic import BaseModel, Field, model_validator +from typing_extensions import Self + +from autogen_agentchat.base import ChatAgent, TerminationCondition +from autogen_agentchat.messages import ( + BaseAgentEvent, + BaseChatMessage, + MessageFactory, + StopMessage, +) +from autogen_agentchat.state import BaseGroupChatManagerState +from autogen_agentchat.teams import BaseGroupChat + +from ..._group_chat._base_group_chat_manager import BaseGroupChatManager +from ..._group_chat._events import GroupChatTermination + +_DIGRAPH_STOP_MESSAGE = "Digraph execution is complete" + + +class DiGraphEdge(BaseModel): + """Represents a directed edge in a :class:`DiGraph`, with an optional execution condition. + + .. warning:: + + This is an experimental feature, and the API will change in the future releases. + + .. warning:: + + If the condition is a callable, it will not be serialized in the model. + + """ + + target: str # Target node name + condition: Union[str, Callable[[BaseChatMessage], bool], None] = Field(default=None) + """(Experimental) Condition to execute this edge. + If None, the edge is unconditional. + If a string, the edge is conditional on the presence of that string in the last agent chat message. + If a callable, the edge is conditional on the callable returning True when given the last message. + """ + + # Using Field to exclude the condition in serialization if it's a callable + condition_function: Callable[[BaseChatMessage], bool] | None = Field(default=None, exclude=True) + activation_group: str = Field(default="") + """Group identifier for forward dependencies. + + When multiple edges point to the same target node, they are grouped by this field. + This allows distinguishing between different cycles or dependency patterns. + + Example: In a graph containing a cycle like A->B->C->B, the two edges pointing to B (A->B and C->B) + can be in different activation groups to control how B is activated. + Defaults to the target node name if not specified. + """ + activation_condition: Literal["all", "any"] = "all" + """Determines how forward dependencies within the same activation_group are evaluated. + + - "all": All edges in this activation group must be satisfied before the target node can execute + - "any": Any single edge in this activation group being satisfied allows the target node to execute + + This is used to handle complex dependency patterns in cyclic graphs where multiple + paths can lead to the same target node. + """ + + @model_validator(mode="after") + def _validate_condition(self) -> "DiGraphEdge": + # Store callable in a separate field and set condition to None for serialization + if callable(self.condition): + self.condition_function = self.condition + # For serialization purposes, we'll set the condition to None + # when storing as a pydantic model/dict + object.__setattr__(self, "condition", None) + + # Set activation_group to target if not already set + if not self.activation_group: + self.activation_group = self.target + + return self + + def check_condition(self, message: BaseChatMessage) -> bool: + """Check if the edge condition is satisfied for the given message. + + Args: + message: The message to check the condition against. + + Returns: + True if condition is satisfied (None condition always returns True), + False otherwise. + """ + if self.condition_function is not None: + return self.condition_function(message) + elif isinstance(self.condition, str): + # If it's a string, check if the string is in the message content + return self.condition in message.to_model_text() + return True # None condition is always satisfied + + +class DiGraphNode(BaseModel): + """Represents a node (agent) in a :class:`DiGraph`, with its outgoing edges and activation type. + + .. warning:: + + This is an experimental feature, and the API will change in the future releases. + + """ + + name: str # Agent's name + edges: List[DiGraphEdge] = [] # Outgoing edges + activation: Literal["all", "any"] = "all" + + +class DiGraph(BaseModel): + """Defines a directed graph structure with nodes and edges. + :class:`GraphFlow` uses this to determine execution order and conditions. + + .. warning:: + + This is an experimental feature, and the API will change in the future releases. + + """ + + nodes: Dict[str, DiGraphNode] # Node name → DiGraphNode mapping + default_start_node: str | None = None # Default start node name + _has_cycles: bool | None = None # Cyclic graph flag + + def get_parents(self) -> Dict[str, List[str]]: + """Compute a mapping of each node to its parent nodes.""" + parents: Dict[str, List[str]] = {node: [] for node in self.nodes} + for node in self.nodes.values(): + for edge in node.edges: + parents[edge.target].append(node.name) + return parents + + def get_start_nodes(self) -> Set[str]: + """Return the nodes that have no incoming edges (entry points).""" + if self.default_start_node: + return {self.default_start_node} + + parents = self.get_parents() + return set([node_name for node_name, parent_list in parents.items() if not parent_list]) + + def get_leaf_nodes(self) -> Set[str]: + """Return nodes that have no outgoing edges (final output nodes).""" + return set([name for name, node in self.nodes.items() if not node.edges]) + + def has_cycles_with_exit(self) -> bool: + """ + Check if the graph has any cycles and validate that each cycle has at least one conditional edge. + + Returns: + bool: True if there is at least one cycle and all cycles have an exit condition. + False if there are no cycles. + + Raises: + ValueError: If there is a cycle without any conditional edge. + """ + visited: Set[str] = set() + rec_stack: Set[str] = set() + path: List[str] = [] + + def dfs(node_name: str) -> bool: + visited.add(node_name) + rec_stack.add(node_name) + path.append(node_name) + + for edge in self.nodes[node_name].edges: + target = edge.target + if target not in visited: + if dfs(target): + return True + elif target in rec_stack: + # Found a cycle → extract the cycle + cycle_start_index = path.index(target) + cycle_nodes = path[cycle_start_index:] + cycle_edges: List[DiGraphEdge] = [] + for n in cycle_nodes: + cycle_edges.extend(self.nodes[n].edges) + if all(edge.condition is None and edge.condition_function is None for edge in cycle_edges): + raise ValueError( + f"Cycle detected without exit condition: {' -> '.join(cycle_nodes + cycle_nodes[:1])}" + ) + return True # Found cycle, but it has an exit condition + + rec_stack.remove(node_name) + path.pop() + return False + + has_cycle = False + for node in self.nodes: + if node not in visited: + if dfs(node): + has_cycle = True + + return has_cycle + + def get_has_cycles(self) -> bool: + """Indicates if the graph has at least one cycle (with valid exit conditions).""" + if self._has_cycles is None: + self._has_cycles = self.has_cycles_with_exit() + + return self._has_cycles + + def graph_validate(self) -> None: + """Validate graph structure and execution rules.""" + if not self.nodes: + raise ValueError("Graph has no nodes.") + + if not self.get_start_nodes(): + raise ValueError("Graph must have at least one start node") + + if not self.get_leaf_nodes(): + raise ValueError("Graph must have at least one leaf node") + + # Outgoing edge condition validation (per node) + for node in self.nodes.values(): + # Check that if a node has an outgoing conditional edge, then all outgoing edges are conditional + has_condition = any( + edge.condition is not None or edge.condition_function is not None for edge in node.edges + ) + has_unconditioned = any(edge.condition is None and edge.condition_function is None for edge in node.edges) + if has_condition and has_unconditioned: + raise ValueError(f"Node '{node.name}' has a mix of conditional and unconditional edges.") + + # Validate activation conditions across all edges in the graph + self._validate_activation_conditions() + + self._has_cycles = self.has_cycles_with_exit() + + def _validate_activation_conditions(self) -> None: + """Validate that all edges pointing to the same target node have consistent activation_condition values. + + Raises: + ValueError: If edges pointing to the same target have different activation_condition values + """ + target_activation_conditions: Dict[str, Dict[str, str]] = {} # target_node -> {activation_group -> condition} + + for node in self.nodes.values(): + for edge in node.edges: + target = edge.target # The target node this edge points to + activation_group = edge.activation_group + + if target not in target_activation_conditions: + target_activation_conditions[target] = {} + + if activation_group in target_activation_conditions[target]: + if target_activation_conditions[target][activation_group] != edge.activation_condition: + # Find the source node that has the conflicting condition + conflicting_source = self._find_edge_source_by_target_and_group( + target, activation_group, target_activation_conditions[target][activation_group] + ) + raise ValueError( + f"Conflicting activation conditions for target '{target}' group '{activation_group}': " + f"'{target_activation_conditions[target][activation_group]}' (from node '{conflicting_source}') " + f"and '{edge.activation_condition}' (from node '{node.name}')" + ) + else: + target_activation_conditions[target][activation_group] = edge.activation_condition + + def _find_edge_source_by_target_and_group( + self, target: str, activation_group: str, activation_condition: str + ) -> str: + """Find the source node that has an edge pointing to the given target with the given activation_group and activation_condition.""" + for node_name, node in self.nodes.items(): + for edge in node.edges: + if ( + edge.target == target + and edge.activation_group == activation_group + and edge.activation_condition == activation_condition + ): + return node_name + return "unknown" + + def get_remaining_map(self) -> Dict[str, Dict[str, int]]: + """Get the remaining map that tracks how many edges point to each target node with each activation group. + + Returns: + Dictionary mapping target nodes to their activation groups and remaining counts + """ + + remaining_map: Dict[str, Dict[str, int]] = {} + + for node in self.nodes.values(): + for edge in node.edges: + target = edge.target + activation_group = edge.activation_group + + if target not in remaining_map: + remaining_map[target] = {} + + if activation_group not in remaining_map[target]: + remaining_map[target][activation_group] = 0 + + remaining_map[target][activation_group] += 1 + + return remaining_map + + +class GraphFlowManagerState(BaseGroupChatManagerState): + """Tracks active execution state for DAG-based execution.""" + + active_nodes: List[str] = [] # Currently executing nodes + type: str = "GraphManagerState" + + +class GraphFlowManager(BaseGroupChatManager): + """Manages execution of agents using a Directed Graph execution model.""" + + def __init__( + self, + name: str, + group_topic_type: str, + output_topic_type: str, + participant_topic_types: List[str], + participant_names: List[str], + participant_descriptions: List[str], + output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination], + termination_condition: TerminationCondition | None, + max_turns: int | None, + message_factory: MessageFactory, + graph: DiGraph, + ) -> None: + """Initialize the graph-based execution manager.""" + super().__init__( + name=name, + group_topic_type=group_topic_type, + output_topic_type=output_topic_type, + participant_topic_types=participant_topic_types, + participant_names=participant_names, + participant_descriptions=participant_descriptions, + output_message_queue=output_message_queue, + termination_condition=termination_condition, + max_turns=max_turns, + message_factory=message_factory, + ) + graph.graph_validate() + if graph.get_has_cycles() and self._termination_condition is None and self._max_turns is None: + raise ValueError("A termination condition is required for cyclic graphs without a maximum turn limit.") + self._graph = graph + # Lookup table for incoming edges for each node. + self._parents = graph.get_parents() + # Lookup table for outgoing edges for each node. + self._edges: Dict[str, List[DiGraphEdge]] = {n: node.edges for n, node in graph.nodes.items()} + + # Build activation and enqueued_any lookup tables by collecting all edges and grouping by target node + self._build_lookup_tables(graph) + + # Track which activation groups were triggered for each node + self._triggered_activation_groups: Dict[str, Set[str]] = {} + # === Mutable states for the graph execution === + # Count the number of remaining parents to activate each node. + self._remaining: Dict[str, Counter[str]] = { + target: Counter(groups) for target, groups in graph.get_remaining_map().items() + } + # cache for remaining + self._origin_remaining: Dict[str, Dict[str, int]] = { + target: Counter(groups) for target, groups in self._remaining.items() + } + + # Ready queue for nodes that are ready to execute, starting with the start nodes. + self._ready: Deque[str] = deque([n for n in graph.get_start_nodes()]) + + def _build_lookup_tables(self, graph: DiGraph) -> None: + """Build activation and enqueued_any lookup tables by collecting all edges and grouping by target node. + + Args: + graph: The directed graph + """ + self._activation: Dict[str, Dict[str, Literal["any", "all"]]] = {} + self._enqueued_any: Dict[str, Dict[str, bool]] = {} + + for node in graph.nodes.values(): + for edge in node.edges: + target = edge.target + activation_group = edge.activation_group + + # Build activation lookup + if target not in self._activation: + self._activation[target] = {} + if activation_group not in self._activation[target]: + self._activation[target][activation_group] = edge.activation_condition + + # Build enqueued_any lookup + if target not in self._enqueued_any: + self._enqueued_any[target] = {} + if activation_group not in self._enqueued_any[target]: + self._enqueued_any[target][activation_group] = False + + async def update_message_thread(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> None: + await super().update_message_thread(messages) + + # Find the node that ran in the current turn. + message = messages[-1] + if message.source not in self._graph.nodes: + # Ignore messages from sources outside of the graph. + return + assert isinstance(message, BaseChatMessage) + source = message.source + + # Propagate the update to the children of the node. + for edge in self._edges[source]: + # Use the new check_condition method that handles both string and callable conditions + if not edge.check_condition(message): + continue + + target = edge.target + activation_group = edge.activation_group + + if self._activation[target][activation_group] == "all": + self._remaining[target][activation_group] -= 1 + if self._remaining[target][activation_group] == 0: + # If all parents are done, add to the ready queue. + self._ready.append(target) + # Track which activation group was triggered + self._save_triggered_activation_group(target, activation_group) + else: + # If activation is any, add to the ready queue if not already enqueued. + if not self._enqueued_any[target][activation_group]: + self._ready.append(target) + self._enqueued_any[target][activation_group] = True + # Track which activation group was triggered + self._save_triggered_activation_group(target, activation_group) + + def _save_triggered_activation_group(self, target: str, activation_group: str) -> None: + """Save which activation group was triggered for a target node. + + Args: + target: The target node that was triggered + activation_group: The activation group that caused the trigger + """ + if target not in self._triggered_activation_groups: + self._triggered_activation_groups[target] = set() + self._triggered_activation_groups[target].add(activation_group) + + def _reset_triggered_activation_groups(self, speaker: str) -> None: + """Reset the bookkeeping for the specific activation groups that were triggered for a speaker. + + Args: + speaker: The speaker node to reset activation groups for + """ + if speaker not in self._triggered_activation_groups: + return + + for activation_group in self._triggered_activation_groups[speaker]: + if self._activation[speaker][activation_group] == "any": + self._enqueued_any[speaker][activation_group] = False + else: + # Reset the remaining count for this activation group using the graph's original count + if speaker in self._remaining and activation_group in self._remaining[speaker]: + self._remaining[speaker][activation_group] = self._origin_remaining[speaker][activation_group] + + # Clear the triggered activation groups for this speaker + self._triggered_activation_groups[speaker].clear() + + async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str]: + # Drain the ready queue for the next set of speakers. + speakers: List[str] = [] + while self._ready: + speaker = self._ready.popleft() + speakers.append(speaker) + + # Reset the bookkeeping for the specific activation groups that were triggered + self._reset_triggered_activation_groups(speaker) + + return speakers + + async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None: + pass + + async def _apply_termination_condition( + self, delta: Sequence[BaseAgentEvent | BaseChatMessage], increment_turn_count: bool = False + ) -> bool: + """Apply termination condition including graph-specific completion logic. + + First checks if graph execution is complete, then checks standard termination conditions. + + Args: + delta: The message delta to check termination conditions against + increment_turn_count: Whether to increment the turn count + + Returns: + True if the conversation should be terminated, False otherwise + """ + # Check if the graph execution is complete (no ready speakers) - prioritize this check + if not self._ready: + stop_message = StopMessage( + content=_DIGRAPH_STOP_MESSAGE, + source=self._name, + ) + # Reset the execution state when the graph has naturally completed + self._reset_execution_state() + # Reset the termination conditions and turn count. + if self._termination_condition is not None: + await self._termination_condition.reset() + self._current_turn = 0 + # Signal termination to the caller of the team. + await self._signal_termination(stop_message) + return True + + # Apply the standard termination conditions from the base class + return await super()._apply_termination_condition(delta, increment_turn_count) + + def _reset_execution_state(self) -> None: + """Reset the graph execution state to the initial state.""" + self._remaining = {target: Counter(groups) for target, groups in self._graph.get_remaining_map().items()} + self._enqueued_any = {n: {g: False for g in self._enqueued_any[n]} for n in self._enqueued_any} + self._ready = deque([n for n in self._graph.get_start_nodes()]) + + async def save_state(self) -> Mapping[str, Any]: + """Save the execution state.""" + state = { + "message_thread": [message.dump() for message in self._message_thread], + "current_turn": self._current_turn, + "remaining": {target: dict(counter) for target, counter in self._remaining.items()}, + "enqueued_any": dict(self._enqueued_any), + "ready": list(self._ready), + } + return state + + async def load_state(self, state: Mapping[str, Any]) -> None: + """Restore execution state from saved data.""" + self._message_thread = [self._message_factory.create(msg) for msg in state["message_thread"]] + self._current_turn = state["current_turn"] + self._remaining = {target: Counter(groups) for target, groups in state["remaining"].items()} + self._enqueued_any = state["enqueued_any"] + self._ready = deque(state["ready"]) + + async def reset(self) -> None: + """Reset execution state to the start of the graph.""" + self._current_turn = 0 + self._message_thread.clear() + if self._termination_condition: + await self._termination_condition.reset() + self._reset_execution_state() + + +class GraphFlowConfig(BaseModel): + """The declarative configuration for GraphFlow.""" + + name: str | None = None + description: str | None = None + participants: List[ComponentModel] + termination_condition: ComponentModel | None = None + max_turns: int | None = None + graph: DiGraph # The execution graph for agents + + +class GraphFlow(BaseGroupChat, Component[GraphFlowConfig]): + """A team that runs a group chat following a Directed Graph execution pattern. + + .. warning:: + + This is an experimental feature, and the API will change in the future releases. + + This group chat executes agents based on a directed graph (:class:`DiGraph`) structure, + allowing complex workflows such as sequential execution, parallel fan-out, + conditional branching, join patterns, and loops with explicit exit conditions. + + The execution order is determined by the edges defined in the `DiGraph`. Each node + in the graph corresponds to an agent, and edges define the flow of messages between agents. + Nodes can be configured to activate when: + + - **All** parent nodes have completed (activation="all") → default + - **Any** parent node completes (activation="any") + + Conditional branching is supported using edge conditions, where the next agent(s) are selected + based on content in the chat history. Loops are permitted as long as there is a condition + that eventually exits the loop. + + .. note:: + + Use the :class:`DiGraphBuilder` class to create a :class:`DiGraph` easily. It provides a fluent API + for adding nodes and edges, setting entry points, and validating the graph structure. + See the :class:`DiGraphBuilder` documentation for more details. + The :class:`GraphFlow` class is designed to be used with the :class:`DiGraphBuilder` for creating complex workflows. + + .. warning:: + + When using callable conditions in edges, they will not be serialized + when calling :meth:`dump_component`. This will be addressed in future releases. + + + Args: + participants (List[ChatAgent]): The participants in the group chat. + termination_condition (TerminationCondition, optional): Termination condition for the chat. + max_turns (int, optional): Maximum number of turns before forcing termination. + graph (DiGraph): Directed execution graph defining node flow and conditions. + + Raises: + ValueError: If participant names are not unique, or if graph validation fails (e.g., cycles without exit). + + Examples: + + **Sequential Flow: A → B → C** + + .. code-block:: python + + import asyncio + + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.conditions import MaxMessageTermination + from autogen_agentchat.teams import DiGraphBuilder, GraphFlow + from autogen_ext.models.openai import OpenAIChatCompletionClient + + + async def main(): + # Initialize agents with OpenAI model clients. + model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano") + agent_a = AssistantAgent("A", model_client=model_client, system_message="You are a helpful assistant.") + agent_b = AssistantAgent("B", model_client=model_client, system_message="Translate input to Chinese.") + agent_c = AssistantAgent("C", model_client=model_client, system_message="Translate input to English.") + + # Create a directed graph with sequential flow A -> B -> C. + builder = DiGraphBuilder() + builder.add_node(agent_a).add_node(agent_b).add_node(agent_c) + builder.add_edge(agent_a, agent_b).add_edge(agent_b, agent_c) + graph = builder.build() + + # Create a GraphFlow team with the directed graph. + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + termination_condition=MaxMessageTermination(5), + ) + + # Run the team and print the events. + async for event in team.run_stream(task="Write a short story about a cat."): + print(event) + + + asyncio.run(main()) + + **Parallel Fan-out: A → (B, C)** + + .. code-block:: python + + import asyncio + + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.conditions import MaxMessageTermination + from autogen_agentchat.teams import DiGraphBuilder, GraphFlow + from autogen_ext.models.openai import OpenAIChatCompletionClient + + + async def main(): + # Initialize agents with OpenAI model clients. + model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano") + agent_a = AssistantAgent("A", model_client=model_client, system_message="You are a helpful assistant.") + agent_b = AssistantAgent("B", model_client=model_client, system_message="Translate input to Chinese.") + agent_c = AssistantAgent("C", model_client=model_client, system_message="Translate input to Japanese.") + + # Create a directed graph with fan-out flow A -> (B, C). + builder = DiGraphBuilder() + builder.add_node(agent_a).add_node(agent_b).add_node(agent_c) + builder.add_edge(agent_a, agent_b).add_edge(agent_a, agent_c) + graph = builder.build() + + # Create a GraphFlow team with the directed graph. + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + termination_condition=MaxMessageTermination(5), + ) + + # Run the team and print the events. + async for event in team.run_stream(task="Write a short story about a cat."): + print(event) + + + asyncio.run(main()) + + **Conditional Branching: A → B (if 'yes') or C (otherwise)** + + .. code-block:: python + + import asyncio + + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.conditions import MaxMessageTermination + from autogen_agentchat.teams import DiGraphBuilder, GraphFlow + from autogen_ext.models.openai import OpenAIChatCompletionClient + + + async def main(): + # Initialize agents with OpenAI model clients. + model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano") + agent_a = AssistantAgent( + "A", + model_client=model_client, + system_message="Detect if the input is in Chinese. If it is, say 'yes', else say 'no', and nothing else.", + ) + agent_b = AssistantAgent("B", model_client=model_client, system_message="Translate input to English.") + agent_c = AssistantAgent("C", model_client=model_client, system_message="Translate input to Chinese.") + + # Create a directed graph with conditional branching flow A -> B ("yes"), A -> C (otherwise). + builder = DiGraphBuilder() + builder.add_node(agent_a).add_node(agent_b).add_node(agent_c) + # Create conditions as callables that check the message content. + builder.add_edge(agent_a, agent_b, condition=lambda msg: "yes" in msg.to_model_text()) + builder.add_edge(agent_a, agent_c, condition=lambda msg: "yes" not in msg.to_model_text()) + graph = builder.build() + + # Create a GraphFlow team with the directed graph. + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + termination_condition=MaxMessageTermination(5), + ) + + # Run the team and print the events. + async for event in team.run_stream(task="AutoGen is a framework for building AI agents."): + print(event) + + + asyncio.run(main()) + + **Loop with exit condition: A → B → C (if 'APPROVE') or A (otherwise)** + + .. code-block:: python + + import asyncio + + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.conditions import MaxMessageTermination + from autogen_agentchat.teams import DiGraphBuilder, GraphFlow + from autogen_ext.models.openai import OpenAIChatCompletionClient + + + async def main(): + # Initialize agents with OpenAI model clients. + model_client = OpenAIChatCompletionClient(model="gpt-4.1") + agent_a = AssistantAgent( + "A", + model_client=model_client, + system_message="You are a helpful assistant.", + ) + agent_b = AssistantAgent( + "B", + model_client=model_client, + system_message="Provide feedback on the input, if your feedback has been addressed, " + "say 'APPROVE', otherwise provide a reason for rejection.", + ) + agent_c = AssistantAgent( + "C", model_client=model_client, system_message="Translate the final product to Korean." + ) + + # Create a loop graph with conditional exit: A -> B -> C ("APPROVE"), B -> A (otherwise). + builder = DiGraphBuilder() + builder.add_node(agent_a).add_node(agent_b).add_node(agent_c) + builder.add_edge(agent_a, agent_b) + + # Create conditional edges using strings + builder.add_edge(agent_b, agent_c, condition=lambda msg: "APPROVE" in msg.to_model_text()) + builder.add_edge(agent_b, agent_a, condition=lambda msg: "APPROVE" not in msg.to_model_text()) + + builder.set_entry_point(agent_a) + graph = builder.build() + + # Create a GraphFlow team with the directed graph. + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + termination_condition=MaxMessageTermination(20), # Max 20 messages to avoid infinite loop. + ) + + # Run the team and print the events. + async for event in team.run_stream(task="Write a short poem about AI Agents."): + print(event) + + + asyncio.run(main()) + """ + + component_config_schema = GraphFlowConfig + component_provider_override = "autogen_agentchat.teams.GraphFlow" + + DEFAULT_NAME = "GraphFlow" + DEFAULT_DESCRIPTION = "A team of agents" + + def __init__( + self, + participants: List[ChatAgent], + graph: DiGraph, + *, + name: str | None = None, + description: str | None = None, + termination_condition: TerminationCondition | None = None, + max_turns: int | None = None, + runtime: AgentRuntime | None = None, + custom_message_types: List[type[BaseAgentEvent | BaseChatMessage]] | None = None, + ) -> None: + self._input_participants = participants + self._input_termination_condition = termination_condition + + for participant in participants: + if not isinstance(participant, ChatAgent): + raise TypeError(f"Participant {participant} must be a ChatAgent.") + + # No longer add _StopAgent or StopMessageTermination + # Termination is now handled directly in GraphFlowManager._apply_termination_condition + super().__init__( + name=name or self.DEFAULT_NAME, + description=description or self.DEFAULT_DESCRIPTION, + participants=list(participants), + group_chat_manager_name="GraphManager", + group_chat_manager_class=GraphFlowManager, + termination_condition=termination_condition, + max_turns=max_turns, + runtime=runtime, + custom_message_types=custom_message_types, + ) + self._graph = graph + + def _create_group_chat_manager_factory( + self, + name: str, + group_topic_type: str, + output_topic_type: str, + participant_topic_types: List[str], + participant_names: List[str], + participant_descriptions: List[str], + output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination], + termination_condition: TerminationCondition | None, + max_turns: int | None, + message_factory: MessageFactory, + ) -> Callable[[], GraphFlowManager]: + """Creates the factory method for initializing the DiGraph-based chat manager.""" + + def _factory() -> GraphFlowManager: + return GraphFlowManager( + name=name, + group_topic_type=group_topic_type, + output_topic_type=output_topic_type, + participant_topic_types=participant_topic_types, + participant_names=participant_names, + participant_descriptions=participant_descriptions, + output_message_queue=output_message_queue, + termination_condition=termination_condition, + max_turns=max_turns, + message_factory=message_factory, + graph=self._graph, + ) + + return _factory + + def _to_config(self) -> GraphFlowConfig: + """Converts the instance into a configuration object.""" + participants = [participant.dump_component() for participant in self._input_participants] + termination_condition = ( + self._input_termination_condition.dump_component() if self._input_termination_condition else None + ) + return GraphFlowConfig( + name=self._name, + description=self._description, + participants=participants, + termination_condition=termination_condition, + max_turns=self._max_turns, + graph=self._graph, + ) + + @classmethod + def _from_config(cls, config: GraphFlowConfig) -> Self: + """Reconstructs an instance from a configuration object.""" + participants = [ChatAgent.load_component(participant) for participant in config.participants] + termination_condition = ( + TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None + ) + return cls( + name=config.name, + description=config.description, + participants=participants, + graph=config.graph, + termination_condition=termination_condition, + max_turns=config.max_turns, + ) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/_graph_builder.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/_graph_builder.py new file mode 100644 index 000000000000..e083415db5a9 --- /dev/null +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/_graph_builder.py @@ -0,0 +1,209 @@ +import warnings +from typing import Callable, Dict, Literal, Optional, Union + +from autogen_agentchat.base import ChatAgent +from autogen_agentchat.messages import BaseChatMessage + +from ._digraph_group_chat import DiGraph, DiGraphEdge, DiGraphNode + + +class DiGraphBuilder: + """ + A fluent builder for constructing :class:`DiGraph` execution graphs used in :class:`GraphFlow`. + + .. warning:: + + This is an experimental feature, and the API will change in the future releases. + + This utility provides a convenient way to programmatically build a graph of agent interactions, + including complex execution flows such as: + + - Sequential chains + - Parallel fan-outs + - Conditional branching + - Cyclic loops with safe exits + + Each node in the graph represents an agent. Edges define execution paths between agents, + and can optionally be conditioned on message content using callable functions. + + The builder is compatible with the `Graph` runner and supports both standard and filtered agents. + + Methods: + - add_node(agent, activation): Add an agent node to the graph. + - add_edge(source, target, condition): Connect two nodes optionally with a condition. + - add_conditional_edges(source, condition_to_target): Add multiple conditional edges from a source. + - set_entry_point(agent): Define the default start node (optional). + - build(): Generate a validated `DiGraph`. + - get_participants(): Return the list of added agents. + + Example — Sequential Flow A → B → C: + >>> builder = GraphBuilder() + >>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c) + >>> builder.add_edge(agent_a, agent_b).add_edge(agent_b, agent_c) + >>> team = Graph( + ... participants=builder.get_participants(), + ... graph=builder.build(), + ... termination_condition=MaxMessageTermination(5), + ... ) + + Example — Parallel Fan-out A → (B, C): + >>> builder = GraphBuilder() + >>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c) + >>> builder.add_edge(agent_a, agent_b).add_edge(agent_a, agent_c) + + Example — Conditional Branching A → B or A → C: + >>> builder = GraphBuilder() + >>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c) + >>> # Add conditional edges using keyword check + >>> builder.add_edge(agent_a, agent_b, condition="keyword1") + >>> builder.add_edge(agent_a, agent_c, condition="keyword2") + + + Example — Using Custom String Conditions: + >>> builder = GraphBuilder() + >>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c) + >>> # Add condition strings to check in messages + >>> builder.add_edge(agent_a, agent_b, condition="big") + >>> builder.add_edge(agent_a, agent_c, condition="small") + + Example — Loop: A → B → A or B → C: + >>> builder = GraphBuilder() + >>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c) + >>> builder.add_edge(agent_a, agent_b) + >> # Add a loop back to agent A + >>> builder.add_edge(agent_b, agent_a, condition=lambda msg: "loop" in msg.to_model_text()) + >>> # Add exit condition to break the loop + >>> builder.add_edge(agent_b, agent_c, condition=lambda msg: "loop" not in msg.to_model_text()) + + Example — Loop with multiple paths to the same node: A → B → C → B: + >>> builder = GraphBuilder() + >>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c) + >>> builder.add_edge(agent_a, agent_b) + >>> builder.add_edge(agent_b, agent_c) + >>> builder.add_edge(agent_c, agent_b, activation_group="loop_back") + + Example — Loop with multiple paths to the same node with any activation condition: A → B → (C1, C2) → B → E(exit): + >>> builder = GraphBuilder() + >>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c1).add_node(agent_c2).add_node(agent_e) + >>> builder.add_edge(agent_a, agent_b) + >>> builder.add_edge(agent_b, agent_c1) + >>> builder.add_edge(agent_b, agent_c2) + >>> builder.add_edge(agent_b, agent_e, condition="exit") + >>> builder.add_edge(agent_c1, agent_b, activation_group="loop_back_group", activation_condition="any") + >>> builder.add_edge(agent_c2, agent_b, activation_group="loop_back_group", activation_condition="any") + """ + + def __init__(self) -> None: + self.nodes: Dict[str, DiGraphNode] = {} + self.agents: Dict[str, ChatAgent] = {} + self._default_start_node: Optional[str] = None + + def _get_name(self, obj: Union[str, ChatAgent]) -> str: + return obj if isinstance(obj, str) else obj.name + + def add_node(self, agent: ChatAgent, activation: Literal["all", "any"] = "all") -> "DiGraphBuilder": + """Add a node to the graph and register its agent.""" + name = agent.name + if name not in self.nodes: + self.nodes[name] = DiGraphNode(name=name, edges=[], activation=activation) + self.agents[name] = agent + return self + + def add_edge( + self, + source: Union[str, ChatAgent], + target: Union[str, ChatAgent], + condition: Optional[Union[str, Callable[[BaseChatMessage], bool]]] = None, + activation_group: Optional[str] = None, + activation_condition: Optional[Literal["all", "any"]] = None, + ) -> "DiGraphBuilder": + """Add a directed edge from source to target, optionally with a condition. + + Args: + source: Source node (agent name or agent object) + target: Target node (agent name or agent object) + condition: Optional condition for edge activation. + If string, activates when substring is found in message. + If callable, activates when function returns True for the message. + + Returns: + Self for method chaining + + Raises: + ValueError: If source or target node doesn't exist in the builder + """ + source_name = self._get_name(source) + target_name = self._get_name(target) + + if source_name not in self.nodes: + raise ValueError(f"Source node '{source_name}' must be added before adding an edge.") + if target_name not in self.nodes: + raise ValueError(f"Target node '{target_name}' must be added before adding an edge.") + if activation_group is None: + activation_group = target_name + if activation_condition is None: + activation_condition = "all" + self.nodes[source_name].edges.append( + DiGraphEdge( + target=target_name, + condition=condition, + activation_group=activation_group, + activation_condition=activation_condition, + ) + ) + return self + + def add_conditional_edges( + self, source: Union[str, ChatAgent], condition_to_target: Dict[str, Union[str, ChatAgent]] + ) -> "DiGraphBuilder": + """Add multiple conditional edges from a source node based on keyword checks. + + .. warning:: + + This method interface will be changed in the future to support callable conditions. + Please use `add_edge` if you need to specify custom conditions. + + Args: + source: Source node (agent name or agent object) + condition_to_target: Mapping from condition strings to target nodes + Each key is a keyword that will be checked in the message content + Each value is the target node to activate when condition is met + + For each key (keyword), a lambda will be created that checks + if the keyword is in the message text. + + Returns: + Self for method chaining + """ + + warnings.warn( + "add_conditional_edges will be changed in the future to support callable conditions. " + "For now, please use add_edge if you need to specify custom conditions.", + DeprecationWarning, + stacklevel=2, + ) + + for condition_keyword, target in condition_to_target.items(): + self.add_edge(source, target, condition=condition_keyword) + return self + + def set_entry_point(self, name: Union[str, ChatAgent]) -> "DiGraphBuilder": + """Set the default start node of the graph.""" + node_name = self._get_name(name) + if node_name not in self.nodes: + raise ValueError(f"Start node '{node_name}' must be added before setting as entry point.") + self._default_start_node = node_name + return self + + def build(self) -> DiGraph: + """Build and validate the DiGraph.""" + graph = DiGraph( + nodes=self.nodes, + default_start_node=self._default_start_node, + ) + graph.graph_validate() + return graph + + def get_participants(self) -> list[ChatAgent]: + """Return the list of agents in the builder, in insertion order.""" + return list(self.agents.values()) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py index d391f7b62ff8..e5fd0e85a862 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py @@ -9,7 +9,7 @@ from .... import EVENT_LOGGER_NAME, TRACE_LOGGER_NAME from ....base import ChatAgent, TerminationCondition -from ....messages import AgentEvent, ChatMessage +from ....messages import BaseAgentEvent, BaseChatMessage, MessageFactory from .._base_group_chat import BaseGroupChat from .._events import GroupChatTermination from ._magentic_one_orchestrator import MagenticOneOrchestrator @@ -22,12 +22,15 @@ class MagenticOneGroupChatConfig(BaseModel): """The declarative configuration for a MagenticOneGroupChat.""" + name: str | None = None + description: str | None = None participants: List[ComponentModel] model_client: ComponentModel termination_condition: ComponentModel | None = None max_turns: int | None = None max_stalls: int final_answer_prompt: str + emit_team_events: bool = False class MagenticOneGroupChat(BaseGroupChat, Component[MagenticOneGroupChatConfig]): @@ -38,6 +41,9 @@ class MagenticOneGroupChat(BaseGroupChat, Component[MagenticOneGroupChatConfig]) The orchestrator is based on the Magentic-One architecture, which is a generalist multi-agent system for solving complex tasks (see references below). + Unlike :class:`~autogen_agentchat.teams.RoundRobinGroupChat` and :class:`~autogen_agentchat.teams.SelectorGroupChat`, + the MagenticOneGroupChat does not support using team as participant. + Args: participants (List[ChatAgent]): The participants in the group chat. model_client (ChatCompletionClient): The model client used for generating responses. @@ -46,6 +52,10 @@ class MagenticOneGroupChat(BaseGroupChat, Component[MagenticOneGroupChatConfig]) max_turns (int, optional): The maximum number of turns in the group chat before stopping. Defaults to 20. max_stalls (int, optional): The maximum number of stalls allowed before re-planning. Defaults to 3. final_answer_prompt (str, optional): The LLM prompt used to generate the final answer or response from the team's transcript. A default (sensible for GPT-4o class models) is provided. + custom_message_types (List[type[BaseAgentEvent | BaseChatMessage]], optional): A list of custom message types that will be used in the group chat. + If you are using custom message types or your agents produces custom message types, you need to specify them here. + Make sure your custom message types are subclasses of :class:`~autogen_agentchat.messages.BaseAgentEvent` or :class:`~autogen_agentchat.messages.BaseChatMessage`. + emit_team_events (bool, optional): Whether to emit team events through :meth:`BaseGroupChat.run_stream`. Defaults to False. Raises: ValueError: In orchestration logic if progress ledger does not have required keys or if next speaker is not valid. @@ -93,24 +103,38 @@ async def main() -> None: component_config_schema = MagenticOneGroupChatConfig component_provider_override = "autogen_agentchat.teams.MagenticOneGroupChat" + DEFAULT_NAME = "MagenticOneGroupChat" + DEFAULT_DESCRIPTION = "A team of agents." + def __init__( self, participants: List[ChatAgent], model_client: ChatCompletionClient, *, + name: str | None = None, + description: str | None = None, termination_condition: TerminationCondition | None = None, max_turns: int | None = 20, runtime: AgentRuntime | None = None, max_stalls: int = 3, final_answer_prompt: str = ORCHESTRATOR_FINAL_ANSWER_PROMPT, + custom_message_types: List[type[BaseAgentEvent | BaseChatMessage]] | None = None, + emit_team_events: bool = False, ): + for participant in participants: + if not isinstance(participant, ChatAgent): + raise TypeError(f"Participant {participant} must be a ChatAgent.") super().__init__( - participants, + name=name or self.DEFAULT_NAME, + description=description or self.DEFAULT_DESCRIPTION, + participants=list(participants), group_chat_manager_name="MagenticOneOrchestrator", group_chat_manager_class=MagenticOneOrchestrator, termination_condition=termination_condition, max_turns=max_turns, runtime=runtime, + custom_message_types=custom_message_types, + emit_team_events=emit_team_events, ) # Validate the participants. @@ -128,9 +152,10 @@ def _create_group_chat_manager_factory( participant_topic_types: List[str], participant_names: List[str], participant_descriptions: List[str], - output_message_queue: asyncio.Queue[AgentEvent | ChatMessage | GroupChatTermination], + output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination], termination_condition: TerminationCondition | None, max_turns: int | None, + message_factory: MessageFactory, ) -> Callable[[], MagenticOneOrchestrator]: return lambda: MagenticOneOrchestrator( name, @@ -140,23 +165,28 @@ def _create_group_chat_manager_factory( participant_names, participant_descriptions, max_turns, + message_factory, self._model_client, self._max_stalls, self._final_answer_prompt, output_message_queue, termination_condition, + self._emit_team_events, ) def _to_config(self) -> MagenticOneGroupChatConfig: participants = [participant.dump_component() for participant in self._participants] termination_condition = self._termination_condition.dump_component() if self._termination_condition else None return MagenticOneGroupChatConfig( + name=self.name, + description=self.description, participants=participants, model_client=self._model_client.dump_component(), termination_condition=termination_condition, max_turns=self._max_turns, max_stalls=self._max_stalls, final_answer_prompt=self._final_answer_prompt, + emit_team_events=self._emit_team_events, ) @classmethod @@ -167,10 +197,13 @@ def _from_config(cls, config: MagenticOneGroupChatConfig) -> Self: TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None ) return cls( - participants, - model_client, + participants=participants, + name=config.name, + description=config.description, + model_client=model_client, termination_condition=termination_condition, max_turns=config.max_turns, max_stalls=config.max_stalls, final_answer_prompt=config.final_answer_prompt, + emit_team_events=config.emit_team_events, ) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_orchestrator.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_orchestrator.py index bfef2b4ed184..176789257ba7 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_orchestrator.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_orchestrator.py @@ -2,7 +2,7 @@ import json import logging import re -from typing import Any, Dict, List, Mapping +from typing import Any, Dict, List, Mapping, Sequence from autogen_core import AgentId, CancellationToken, DefaultTopicId, MessageContext, event, rpc from autogen_core.models import ( @@ -11,14 +11,17 @@ LLMMessage, UserMessage, ) +from autogen_core.utils import extract_json_from_str from .... import TRACE_LOGGER_NAME from ....base import Response, TerminationCondition from ....messages import ( - AgentEvent, - ChatMessage, + BaseAgentEvent, + BaseChatMessage, HandoffMessage, + MessageFactory, MultiModalMessage, + SelectSpeakerEvent, StopMessage, TextMessage, ToolCallExecutionEvent, @@ -26,7 +29,7 @@ ToolCallSummaryMessage, ) from ....state import MagenticOneOrchestratorState -from ....utils import content_to_str, remove_images +from ....utils import remove_images from .._base_group_chat_manager import BaseGroupChatManager from .._events import ( GroupChatAgentResponse, @@ -34,7 +37,9 @@ GroupChatRequestPublish, GroupChatReset, GroupChatStart, + GroupChatTeamResponse, GroupChatTermination, + SerializableException, ) from ._prompts import ( ORCHESTRATOR_FINAL_ANSWER_PROMPT, @@ -44,6 +49,7 @@ ORCHESTRATOR_TASK_LEDGER_FULL_PROMPT, ORCHESTRATOR_TASK_LEDGER_PLAN_PROMPT, ORCHESTRATOR_TASK_LEDGER_PLAN_UPDATE_PROMPT, + LedgerEntry, ) trace_logger = logging.getLogger(TRACE_LOGGER_NAME) @@ -61,11 +67,13 @@ def __init__( participant_names: List[str], participant_descriptions: List[str], max_turns: int | None, + message_factory: MessageFactory, model_client: ChatCompletionClient, max_stalls: int, final_answer_prompt: str, - output_message_queue: asyncio.Queue[AgentEvent | ChatMessage | GroupChatTermination], + output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination], termination_condition: TerminationCondition | None, + emit_team_events: bool, ): super().__init__( name, @@ -77,6 +85,8 @@ def __init__( output_message_queue, termination_condition, max_turns, + message_factory, + emit_team_events=emit_team_events, ) self._model_client = model_client self._max_stalls = max_stalls @@ -147,7 +157,7 @@ async def handle_start(self, message: GroupChatStart, ctx: MessageContext) -> No # Create the initial task ledger ################################# # Combine all message contents for task - self._task = " ".join([content_to_str(msg.content) for msg in message.messages]) + self._task = " ".join([msg.to_model_text() for msg in message.messages]) planning_conversation: List[LLMMessage] = [] # 1. GATHER FACTS @@ -180,30 +190,41 @@ async def handle_start(self, message: GroupChatStart, ctx: MessageContext) -> No await self._reenter_outer_loop(ctx.cancellation_token) @event - async def handle_agent_response(self, message: GroupChatAgentResponse, ctx: MessageContext) -> None: # type: ignore - delta: List[AgentEvent | ChatMessage] = [] - if message.agent_response.inner_messages is not None: - for inner_message in message.agent_response.inner_messages: - delta.append(inner_message) - self._message_thread.append(message.agent_response.chat_message) - delta.append(message.agent_response.chat_message) - - if self._termination_condition is not None: - stop_message = await self._termination_condition(delta) - if stop_message is not None: - # Reset the termination conditions. - await self._termination_condition.reset() - # Signal termination. - await self._signal_termination(stop_message) - return - await self._orchestrate_step(ctx.cancellation_token) - - async def validate_group_state(self, messages: List[ChatMessage] | None) -> None: + async def handle_agent_response( # type: ignore + self, message: GroupChatAgentResponse | GroupChatTeamResponse, ctx: MessageContext + ) -> None: # type: ignore + try: + if not isinstance(message, GroupChatAgentResponse): + raise RuntimeError("MagenticOneOrchestrator does not support GroupChatTeamResponse messages.") + delta: List[BaseAgentEvent | BaseChatMessage] = [] + if message.response.inner_messages is not None: + for inner_message in message.response.inner_messages: + delta.append(inner_message) + await self.update_message_thread([message.response.chat_message]) + delta.append(message.response.chat_message) + + if self._termination_condition is not None: + stop_message = await self._termination_condition(delta) + if stop_message is not None: + # Reset the termination conditions. + await self._termination_condition.reset() + # Signal termination. + await self._signal_termination(stop_message) + return + + await self._orchestrate_step(ctx.cancellation_token) + except Exception as e: + error = SerializableException.from_exception(e) + await self._signal_termination_with_error(error) + # Raise the error to the runtime. + raise + + async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None: pass async def save_state(self) -> Mapping[str, Any]: state = MagenticOneOrchestratorState( - message_thread=list(self._message_thread), + message_thread=[msg.dump() for msg in self._message_thread], current_turn=self._current_turn, task=self._task, facts=self._facts, @@ -215,7 +236,7 @@ async def save_state(self) -> Mapping[str, Any]: async def load_state(self, state: Mapping[str, Any]) -> None: orchestrator_state = MagenticOneOrchestratorState.model_validate(state) - self._message_thread = orchestrator_state.message_thread + self._message_thread = [self._message_factory.create(message) for message in orchestrator_state.message_thread] self._current_turn = orchestrator_state.current_turn self._task = orchestrator_state.task self._facts = orchestrator_state.facts @@ -223,9 +244,9 @@ async def load_state(self, state: Mapping[str, Any]) -> None: self._n_rounds = orchestrator_state.n_rounds self._n_stalls = orchestrator_state.n_stalls - async def select_speaker(self, thread: List[AgentEvent | ChatMessage]) -> str: + async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str: """Not used in this orchestrator, we select next speaker in _orchestrate_step.""" - return "" + return [""] async def reset(self) -> None: """Reset the group chat manager.""" @@ -257,7 +278,7 @@ async def _reenter_outer_loop(self, cancellation_token: CancellationToken) -> No ) # Save my copy - self._message_thread.append(ledger_message) + await self.update_message_thread([ledger_message]) # Log it to the output topic. await self.publish_message( @@ -269,7 +290,7 @@ async def _reenter_outer_loop(self, cancellation_token: CancellationToken) -> No # Broadcast await self.publish_message( - GroupChatAgentResponse(agent_response=Response(chat_message=ledger_message)), + GroupChatAgentResponse(response=Response(chat_message=ledger_message), name=self._name), topic_id=DefaultTopicId(type=self._group_topic_type), ) @@ -295,11 +316,27 @@ async def _orchestrate_step(self, cancellation_token: CancellationToken) -> None assert self._max_json_retries > 0 key_error: bool = False for _ in range(self._max_json_retries): - response = await self._model_client.create(self._get_compatible_context(context), json_output=True) + if self._model_client.model_info.get("structured_output", False): + response = await self._model_client.create( + self._get_compatible_context(context), json_output=LedgerEntry + ) + elif self._model_client.model_info.get("json_output", False): + response = await self._model_client.create( + self._get_compatible_context(context), cancellation_token=cancellation_token, json_output=True + ) + else: + response = await self._model_client.create( + self._get_compatible_context(context), cancellation_token=cancellation_token + ) ledger_str = response.content try: assert isinstance(ledger_str, str) - progress_ledger = json.loads(ledger_str) + output_json = extract_json_from_str(ledger_str) + if len(output_json) != 1: + raise ValueError( + f"Progress ledger should contain a single JSON object, but found: {len(progress_ledger)}" + ) + progress_ledger = output_json[0] # If the team consists of a single agent, deterministically set the next speaker if len(self._participant_names) == 1: @@ -370,7 +407,7 @@ async def _orchestrate_step(self, cancellation_token: CancellationToken) -> None # Broadcast the next step message = TextMessage(content=progress_ledger["instruction_or_question"]["answer"], source=self._name) - self._message_thread.append(message) # My copy + await self.update_message_thread([message]) # My copy await self._log_message(f"Next Speaker: {progress_ledger['next_speaker']['answer']}") # Log it to the output topic. @@ -383,7 +420,7 @@ async def _orchestrate_step(self, cancellation_token: CancellationToken) -> None # Broadcast it await self.publish_message( # Broadcast - GroupChatAgentResponse(agent_response=Response(chat_message=message)), + GroupChatAgentResponse(response=Response(chat_message=message), name=self._name), topic_id=DefaultTopicId(type=self._group_topic_type), cancellation_token=cancellation_token, ) @@ -402,6 +439,15 @@ async def _orchestrate_step(self, cancellation_token: CancellationToken) -> None cancellation_token=cancellation_token, ) + # Send the message to the next speaker + if self._emit_team_events: + select_msg = SelectSpeakerEvent(content=[next_speaker], source=self._name) + await self.publish_message( + GroupChatMessage(message=select_msg), + topic_id=DefaultTopicId(type=self._output_topic_type), + ) + await self._output_message_queue.put(select_msg) + async def _update_task_ledger(self, cancellation_token: CancellationToken) -> None: """Update the task ledger (outer loop) with the latest facts and plan.""" context = self._thread_to_context() @@ -443,7 +489,7 @@ async def _prepare_final_answer(self, reason: str, cancellation_token: Cancellat assert isinstance(response.content, str) message = TextMessage(content=response.content, source=self._name) - self._message_thread.append(message) # My copy + await self.update_message_thread([message]) # My copy # Log it to the output topic. await self.publish_message( @@ -455,7 +501,7 @@ async def _prepare_final_answer(self, reason: str, cancellation_token: Cancellat # Broadcast await self.publish_message( - GroupChatAgentResponse(agent_response=Response(chat_message=message)), + GroupChatAgentResponse(response=Response(chat_message=message), name=self._name), topic_id=DefaultTopicId(type=self._group_topic_type), cancellation_token=cancellation_token, ) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_prompts.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_prompts.py index bc3f1b20ec3d..846d06999686 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_prompts.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_prompts.py @@ -1,3 +1,5 @@ +from pydantic import BaseModel + ORCHESTRATOR_SYSTEM_MESSAGE = "" @@ -98,6 +100,24 @@ """ +class LedgerEntryBooleanAnswer(BaseModel): + reason: str + answer: bool + + +class LedgerEntryStringAnswer(BaseModel): + reason: str + answer: str + + +class LedgerEntry(BaseModel): + is_request_satisfied: LedgerEntryBooleanAnswer + is_in_loop: LedgerEntryBooleanAnswer + is_progress_being_made: LedgerEntryBooleanAnswer + next_speaker: LedgerEntryStringAnswer + instruction_or_question: LedgerEntryStringAnswer + + ORCHESTRATOR_TASK_LEDGER_FACTS_UPDATE_PROMPT = """As a reminder, we are working to solve the following task: {task} diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py index 0e630df2d7cb..3f529f0c4474 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py @@ -1,12 +1,12 @@ import asyncio -from typing import Any, Callable, List, Mapping +from typing import Any, Callable, List, Mapping, Sequence from autogen_core import AgentRuntime, Component, ComponentModel from pydantic import BaseModel from typing_extensions import Self -from ...base import ChatAgent, TerminationCondition -from ...messages import AgentEvent, ChatMessage +from ...base import ChatAgent, Team, TerminationCondition +from ...messages import BaseAgentEvent, BaseChatMessage, MessageFactory from ...state import RoundRobinManagerState from ._base_group_chat import BaseGroupChat from ._base_group_chat_manager import BaseGroupChatManager @@ -24,9 +24,11 @@ def __init__( participant_topic_types: List[str], participant_names: List[str], participant_descriptions: List[str], - output_message_queue: asyncio.Queue[AgentEvent | ChatMessage | GroupChatTermination], + output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination], termination_condition: TerminationCondition | None, - max_turns: int | None = None, + max_turns: int | None, + message_factory: MessageFactory, + emit_team_events: bool, ) -> None: super().__init__( name, @@ -38,10 +40,12 @@ def __init__( output_message_queue, termination_condition, max_turns, + message_factory, + emit_team_events, ) self._next_speaker_index = 0 - async def validate_group_state(self, messages: List[ChatMessage] | None) -> None: + async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None: pass async def reset(self) -> None: @@ -53,7 +57,7 @@ async def reset(self) -> None: async def save_state(self) -> Mapping[str, Any]: state = RoundRobinManagerState( - message_thread=list(self._message_thread), + message_thread=[message.dump() for message in self._message_thread], current_turn=self._current_turn, next_speaker_index=self._next_speaker_index, ) @@ -61,12 +65,17 @@ async def save_state(self) -> Mapping[str, Any]: async def load_state(self, state: Mapping[str, Any]) -> None: round_robin_state = RoundRobinManagerState.model_validate(state) - self._message_thread = list(round_robin_state.message_thread) + self._message_thread = [self._message_factory.create(message) for message in round_robin_state.message_thread] self._current_turn = round_robin_state.current_turn self._next_speaker_index = round_robin_state.next_speaker_index - async def select_speaker(self, thread: List[AgentEvent | ChatMessage]) -> str: - """Select a speaker from the participants in a round-robin fashion.""" + async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str: + """Select a speaker from the participants in a round-robin fashion. + + .. note:: + + This method always returns a single speaker. + """ current_speaker_index = self._next_speaker_index self._next_speaker_index = (current_speaker_index + 1) % len(self._participant_names) current_speaker = self._participant_names[current_speaker_index] @@ -76,102 +85,183 @@ async def select_speaker(self, thread: List[AgentEvent | ChatMessage]) -> str: class RoundRobinGroupChatConfig(BaseModel): """The declarative configuration RoundRobinGroupChat.""" + name: str | None = None + description: str | None = None participants: List[ComponentModel] termination_condition: ComponentModel | None = None max_turns: int | None = None + emit_team_events: bool = False class RoundRobinGroupChat(BaseGroupChat, Component[RoundRobinGroupChatConfig]): """A team that runs a group chat with participants taking turns in a round-robin fashion to publish a message to all. + If an :class:`~autogen_agentchat.base.ChatAgent` is a participant, + the :class:`~autogen_agentchat.messages.BaseChatMessage` from the agent response's + :attr:`~autogen_agentchat.base.Response.chat_message` will be published + to other participants in the group chat. + + If a :class:`~autogen_agentchat.base.Team` is a participant, + the :class:`~autogen_agentchat.messages.BaseChatMessage` + from the team result' :attr:`~autogen_agentchat.base.TaskResult.messages` will be published + to other participants in the group chat. + If a single participant is in the team, the participant will be the only speaker. Args: - participants (List[BaseChatAgent]): The participants in the group chat. + participants (List[ChatAgent | Team]): The participants in the group chat. + name (str | None, optional): The name of the group chat, using :attr:`~autogen_agentchat.teams.RoundRobinGroupChat.DEFAULT_NAME` if not provided. + The name is used by a parent team to identify this group chat so it must be unique within the parent team. + description (str | None, optional): The description of the group chat, using :attr:`~autogen_agentchat.teams.RoundRobinGroupChat.DEFAULT_DESCRIPTION` if not provided. termination_condition (TerminationCondition, optional): The termination condition for the group chat. Defaults to None. Without a termination condition, the group chat will run indefinitely. max_turns (int, optional): The maximum number of turns in the group chat before stopping. Defaults to None, meaning no limit. + custom_message_types (List[type[BaseAgentEvent | BaseChatMessage]], optional): A list of custom message types that will be used in the group chat. + If you are using custom message types or your agents produces custom message types, you need to specify them here. + Make sure your custom message types are subclasses of :class:`~autogen_agentchat.messages.BaseAgentEvent` or :class:`~autogen_agentchat.messages.BaseChatMessage`. + emit_team_events (bool, optional): Whether to emit team events through :meth:`BaseGroupChat.run_stream`. Defaults to False. Raises: ValueError: If no participants are provided or if participant names are not unique. Examples: - A team with one participant with tools: + A team with one participant with tools: + + .. code-block:: python + + import asyncio + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.teams import RoundRobinGroupChat + from autogen_agentchat.conditions import TextMentionTermination + from autogen_agentchat.ui import Console + + + async def main() -> None: + model_client = OpenAIChatCompletionClient(model="gpt-4o") + + async def get_weather(location: str) -> str: + return f"The weather in {location} is sunny." + + assistant = AssistantAgent( + "Assistant", + model_client=model_client, + tools=[get_weather], + ) + termination = TextMentionTermination("TERMINATE") + team = RoundRobinGroupChat([assistant], termination_condition=termination) + await Console(team.run_stream(task="What's the weather in New York?")) + + + asyncio.run(main()) - .. code-block:: python + A team with multiple participants: - import asyncio - from autogen_ext.models.openai import OpenAIChatCompletionClient - from autogen_agentchat.agents import AssistantAgent - from autogen_agentchat.teams import RoundRobinGroupChat - from autogen_agentchat.conditions import TextMentionTermination - from autogen_agentchat.ui import Console + .. code-block:: python + import asyncio + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.teams import RoundRobinGroupChat + from autogen_agentchat.conditions import TextMentionTermination + from autogen_agentchat.ui import Console - async def main() -> None: - model_client = OpenAIChatCompletionClient(model="gpt-4o") - async def get_weather(location: str) -> str: - return f"The weather in {location} is sunny." + async def main() -> None: + model_client = OpenAIChatCompletionClient(model="gpt-4o") - assistant = AssistantAgent( - "Assistant", - model_client=model_client, - tools=[get_weather], - ) - termination = TextMentionTermination("TERMINATE") - team = RoundRobinGroupChat([assistant], termination_condition=termination) - await Console(team.run_stream(task="What's the weather in New York?")) + agent1 = AssistantAgent("Assistant1", model_client=model_client) + agent2 = AssistantAgent("Assistant2", model_client=model_client) + termination = TextMentionTermination("TERMINATE") + team = RoundRobinGroupChat([agent1, agent2], termination_condition=termination) + await Console(team.run_stream(task="Tell me some jokes.")) - asyncio.run(main()) + asyncio.run(main()) - A team with multiple participants: + A team of user proxy and a nested team of writer and reviewer agents: - .. code-block:: python + .. code-block:: python - import asyncio - from autogen_ext.models.openai import OpenAIChatCompletionClient - from autogen_agentchat.agents import AssistantAgent - from autogen_agentchat.teams import RoundRobinGroupChat - from autogen_agentchat.conditions import TextMentionTermination - from autogen_agentchat.ui import Console + import asyncio + from autogen_agentchat.agents import UserProxyAgent, AssistantAgent + from autogen_agentchat.conditions import TextMentionTermination, MaxMessageTermination + from autogen_agentchat.teams import RoundRobinGroupChat + from autogen_agentchat.ui import Console + from autogen_ext.models.openai import OpenAIChatCompletionClient - async def main() -> None: - model_client = OpenAIChatCompletionClient(model="gpt-4o") - agent1 = AssistantAgent("Assistant1", model_client=model_client) - agent2 = AssistantAgent("Assistant2", model_client=model_client) - termination = TextMentionTermination("TERMINATE") - team = RoundRobinGroupChat([agent1, agent2], termination_condition=termination) - await Console(team.run_stream(task="Tell me some jokes.")) + async def main() -> None: + model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano") + writer = AssistantAgent( + "writer", model_client=model_client, system_message="You are a writer.", model_client_stream=True + ) - asyncio.run(main()) + reviewer = AssistantAgent( + "reviewer", + model_client=model_client, + system_message="Provide feedback to the input and suggest improvements.", + model_client_stream=True, + ) + + # NOTE: you can skip input by pressing Enter. + user_proxy = UserProxyAgent("user_proxy") + + # Maximum 1 round of review and revision. + inner_termination = MaxMessageTermination(max_messages=4) + + # The outter-loop termination condition that will terminate the team when the user types "exit". + outter_termination = TextMentionTermination("exit", sources=["user_proxy"]) + + team = RoundRobinGroupChat( + [ + # For each turn, the writer writes a summary and the reviewer reviews it. + RoundRobinGroupChat([writer, reviewer], termination_condition=inner_termination), + # The user proxy gets user input once the writer and reviewer have finished their actions. + user_proxy, + ], + termination_condition=outter_termination, + ) + # Start the team and wait for it to terminate. + await Console(team.run_stream(task="Write a short essay about the impact of AI on society.")) + + + asyncio.run(main()) """ component_config_schema = RoundRobinGroupChatConfig component_provider_override = "autogen_agentchat.teams.RoundRobinGroupChat" - # TODO: Add * to the constructor to separate the positional parameters from the kwargs. - # This may be a breaking change so let's wait until a good time to do it. + DEFAULT_NAME = "RoundRobinGroupChat" + DEFAULT_DESCRIPTION = "A team of agents." + def __init__( self, - participants: List[ChatAgent], + participants: List[ChatAgent | Team], + *, + name: str | None = None, + description: str | None = None, termination_condition: TerminationCondition | None = None, max_turns: int | None = None, runtime: AgentRuntime | None = None, + custom_message_types: List[type[BaseAgentEvent | BaseChatMessage]] | None = None, + emit_team_events: bool = False, ) -> None: super().__init__( - participants, + name=name or self.DEFAULT_NAME, + description=description or self.DEFAULT_DESCRIPTION, + participants=participants, group_chat_manager_name="RoundRobinGroupChatManager", group_chat_manager_class=RoundRobinGroupChatManager, termination_condition=termination_condition, max_turns=max_turns, runtime=runtime, + custom_message_types=custom_message_types, + emit_team_events=emit_team_events, ) def _create_group_chat_manager_factory( @@ -182,9 +272,10 @@ def _create_group_chat_manager_factory( participant_topic_types: List[str], participant_names: List[str], participant_descriptions: List[str], - output_message_queue: asyncio.Queue[AgentEvent | ChatMessage | GroupChatTermination], + output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination], termination_condition: TerminationCondition | None, max_turns: int | None, + message_factory: MessageFactory, ) -> Callable[[], RoundRobinGroupChatManager]: def _factory() -> RoundRobinGroupChatManager: return RoundRobinGroupChatManager( @@ -197,6 +288,8 @@ def _factory() -> RoundRobinGroupChatManager: output_message_queue, termination_condition, max_turns, + message_factory, + self._emit_team_events, ) return _factory @@ -205,15 +298,31 @@ def _to_config(self) -> RoundRobinGroupChatConfig: participants = [participant.dump_component() for participant in self._participants] termination_condition = self._termination_condition.dump_component() if self._termination_condition else None return RoundRobinGroupChatConfig( + name=self._name, + description=self._description, participants=participants, termination_condition=termination_condition, max_turns=self._max_turns, + emit_team_events=self._emit_team_events, ) @classmethod def _from_config(cls, config: RoundRobinGroupChatConfig) -> Self: - participants = [ChatAgent.load_component(participant) for participant in config.participants] + participants: List[ChatAgent | Team] = [] + for participant in config.participants: + if participant.component_type == Team.component_type: + participants.append(Team.load_component(participant)) + else: + participants.append(ChatAgent.load_component(participant)) + termination_condition = ( TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None ) - return cls(participants, termination_condition=termination_condition, max_turns=config.max_turns) + return cls( + participants, + name=config.name, + description=config.description, + termination_condition=termination_condition, + max_turns=config.max_turns, + emit_team_events=config.emit_team_events, + ) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py index 6a2231873a0c..480dc6b71641 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py @@ -4,19 +4,32 @@ from inspect import iscoroutinefunction from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Sequence, Union, cast -from autogen_core import AgentRuntime, Component, ComponentModel -from autogen_core.models import AssistantMessage, ChatCompletionClient, ModelFamily, SystemMessage, UserMessage +from autogen_core import AgentRuntime, CancellationToken, Component, ComponentModel +from autogen_core.model_context import ( + ChatCompletionContext, + UnboundedChatCompletionContext, +) +from autogen_core.models import ( + AssistantMessage, + ChatCompletionClient, + CreateResult, + LLMMessage, + ModelFamily, + SystemMessage, + UserMessage, +) from pydantic import BaseModel from typing_extensions import Self from ... import TRACE_LOGGER_NAME -from ...agents import BaseChatAgent -from ...base import ChatAgent, TerminationCondition +from ...base import ChatAgent, Team, TerminationCondition from ...messages import ( - AgentEvent, BaseAgentEvent, - ChatMessage, - MultiModalMessage, + BaseChatMessage, + HandoffMessage, + MessageFactory, + ModelClientStreamingChunkEvent, + SelectorEvent, ) from ...state import SelectorManagerState from ._base_group_chat import BaseGroupChat @@ -25,12 +38,12 @@ trace_logger = logging.getLogger(TRACE_LOGGER_NAME) -SyncSelectorFunc = Callable[[Sequence[AgentEvent | ChatMessage]], str | None] -AsyncSelectorFunc = Callable[[Sequence[AgentEvent | ChatMessage]], Awaitable[str | None]] +SyncSelectorFunc = Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], str | None] +AsyncSelectorFunc = Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[str | None]] SelectorFuncType = Union[SyncSelectorFunc | AsyncSelectorFunc] -SyncCandidateFunc = Callable[[Sequence[AgentEvent | ChatMessage]], List[str]] -AsyncCandidateFunc = Callable[[Sequence[AgentEvent | ChatMessage]], Awaitable[List[str]]] +SyncCandidateFunc = Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], List[str]] +AsyncCandidateFunc = Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[List[str]]] CandidateFuncType = Union[SyncCandidateFunc | AsyncCandidateFunc] @@ -46,15 +59,19 @@ def __init__( participant_topic_types: List[str], participant_names: List[str], participant_descriptions: List[str], - output_message_queue: asyncio.Queue[AgentEvent | ChatMessage | GroupChatTermination], + output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination], termination_condition: TerminationCondition | None, max_turns: int | None, + message_factory: MessageFactory, model_client: ChatCompletionClient, selector_prompt: str, allow_repeated_speaker: bool, selector_func: Optional[SelectorFuncType], max_selector_attempts: int, candidate_func: Optional[CandidateFuncType], + emit_team_events: bool, + model_context: ChatCompletionContext | None, + model_client_streaming: bool = False, ) -> None: super().__init__( name, @@ -66,6 +83,8 @@ def __init__( output_message_queue, termination_condition, max_turns, + message_factory, + emit_team_events, ) self._model_client = model_client self._selector_prompt = selector_prompt @@ -76,20 +95,27 @@ def __init__( self._max_selector_attempts = max_selector_attempts self._candidate_func = candidate_func self._is_candidate_func_async = iscoroutinefunction(self._candidate_func) + self._model_client_streaming = model_client_streaming + if model_context is not None: + self._model_context = model_context + else: + self._model_context = UnboundedChatCompletionContext() + self._cancellation_token = CancellationToken() - async def validate_group_state(self, messages: List[ChatMessage] | None) -> None: + async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None: pass async def reset(self) -> None: self._current_turn = 0 self._message_thread.clear() + await self._model_context.clear() if self._termination_condition is not None: await self._termination_condition.reset() self._previous_speaker = None async def save_state(self) -> Mapping[str, Any]: state = SelectorManagerState( - message_thread=list(self._message_thread), + message_thread=[msg.dump() for msg in self._message_thread], current_turn=self._current_turn, previous_speaker=self._previous_speaker, ) @@ -97,17 +123,42 @@ async def save_state(self) -> Mapping[str, Any]: async def load_state(self, state: Mapping[str, Any]) -> None: selector_state = SelectorManagerState.model_validate(state) - self._message_thread = list(selector_state.message_thread) + self._message_thread = [self._message_factory.create(msg) for msg in selector_state.message_thread] + await self._add_messages_to_context( + self._model_context, [msg for msg in self._message_thread if isinstance(msg, BaseChatMessage)] + ) self._current_turn = selector_state.current_turn self._previous_speaker = selector_state.previous_speaker - async def select_speaker(self, thread: List[AgentEvent | ChatMessage]) -> str: + @staticmethod + async def _add_messages_to_context( + model_context: ChatCompletionContext, + messages: Sequence[BaseChatMessage], + ) -> None: + """ + Add incoming messages to the model context. + """ + for msg in messages: + if isinstance(msg, HandoffMessage): + for llm_msg in msg.context: + await model_context.add_message(llm_msg) + await model_context.add_message(msg.to_model_message()) + + async def update_message_thread(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> None: + self._message_thread.extend(messages) + base_chat_messages = [m for m in messages if isinstance(m, BaseChatMessage)] + await self._add_messages_to_context(self._model_context, base_chat_messages) + + async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str: """Selects the next speaker in a group chat using a ChatCompletion client, with the selector function as override if it returns a speaker name. + .. note:: + + This method always returns a single speaker name. + A key assumption is that the agent type is the same as the topic type, which we use as the agent name. """ - # Use the selector function if provided. if self._selector_func is not None: if self._is_selector_func_async: @@ -123,7 +174,7 @@ async def select_speaker(self, thread: List[AgentEvent | ChatMessage]) -> str: f"Expected one of: {self._participant_names}." ) # Skip the model based selection. - return speaker + return [speaker] # Use the candidate function to filter participants if provided if self._candidate_func is not None: @@ -149,28 +200,6 @@ async def select_speaker(self, thread: List[AgentEvent | ChatMessage]) -> str: assert len(participants) > 0 - # Construct the history of the conversation. - history_messages: List[str] = [] - for msg in thread: - if isinstance(msg, BaseAgentEvent): - # Ignore agent events. - continue - message = f"{msg.source}:" - if isinstance(msg.content, str): - message += f" {msg.content}" - elif isinstance(msg, MultiModalMessage): - for item in msg.content: - if isinstance(item, str): - message += f" {item}" - else: - message += " [Image]" - else: - raise ValueError(f"Unexpected message type in selector: {type(msg)}") - history_messages.append( - message.rstrip() + "\n\n" - ) # Create some consistency for how messages are separated in the transcript - history = "\n".join(history_messages) - # Construct agent roles. # Each agent sould appear on a single line. roles = "" @@ -180,17 +209,34 @@ async def select_speaker(self, thread: List[AgentEvent | ChatMessage]) -> str: # Select the next speaker. if len(participants) > 1: - agent_name = await self._select_speaker(roles, participants, history, self._max_selector_attempts) + agent_name = await self._select_speaker(roles, participants, self._max_selector_attempts) else: agent_name = participants[0] self._previous_speaker = agent_name trace_logger.debug(f"Selected speaker: {agent_name}") - return agent_name + return [agent_name] + + def construct_message_history(self, message_history: List[LLMMessage]) -> str: + # Construct the history of the conversation. + history_messages: List[str] = [] + for msg in message_history: + if isinstance(msg, UserMessage) or isinstance(msg, AssistantMessage): + message = f"{msg.source}: {msg.content}" + history_messages.append( + message.rstrip() + "\n\n" + ) # Create some consistency for how messages are separated in the transcript + + history: str = "\n".join(history_messages) + return history + + async def _select_speaker(self, roles: str, participants: List[str], max_attempts: int) -> str: + model_context_messages = await self._model_context.get_messages() + model_context_history = self.construct_message_history(model_context_messages) - async def _select_speaker(self, roles: str, participants: List[str], history: str, max_attempts: int) -> str: select_speaker_prompt = self._selector_prompt.format( - roles=roles, participants=str(participants), history=history + roles=roles, participants=str(participants), history=model_context_history ) + select_speaker_messages: List[SystemMessage | UserMessage | AssistantMessage] if ModelFamily.is_openai(self._model_client.model_info["family"]): select_speaker_messages = [SystemMessage(content=select_speaker_prompt)] @@ -201,7 +247,26 @@ async def _select_speaker(self, roles: str, participants: List[str], history: st num_attempts = 0 while num_attempts < max_attempts: num_attempts += 1 - response = await self._model_client.create(messages=select_speaker_messages) + if self._model_client_streaming: + chunk: CreateResult | str = "" + async for _chunk in self._model_client.create_stream(messages=select_speaker_messages): + chunk = _chunk + if self._emit_team_events: + if isinstance(chunk, str): + await self._output_message_queue.put( + ModelClientStreamingChunkEvent(content=cast(str, _chunk), source=self._name) + ) + else: + assert isinstance(chunk, CreateResult) + assert isinstance(chunk.content, str) + await self._output_message_queue.put( + SelectorEvent(content=chunk.content, source=self._name) + ) + # The last chunk must be CreateResult. + assert isinstance(chunk, CreateResult) + response = chunk + else: + response = await self._model_client.create(messages=select_speaker_messages) assert isinstance(response.content, str) select_speaker_messages.append(AssistantMessage(content=response.content, source="selector")) # NOTE: we use all participant names to check for mentions, even if the previous speaker is not allowed. @@ -279,6 +344,8 @@ def _mentioned_agents(self, message_content: str, agent_names: List[str]) -> Dic class SelectorGroupChatConfig(BaseModel): """The declarative configuration for SelectorGroupChat.""" + name: str | None = None + description: str | None = None participants: List[ComponentModel] model_client: ComponentModel termination_condition: ComponentModel | None = None @@ -287,36 +354,65 @@ class SelectorGroupChatConfig(BaseModel): allow_repeated_speaker: bool # selector_func: ComponentModel | None max_selector_attempts: int = 3 + emit_team_events: bool = False + model_client_streaming: bool = False + model_context: ComponentModel | None = None class SelectorGroupChat(BaseGroupChat, Component[SelectorGroupChatConfig]): """A group chat team that have participants takes turn to publish a message to all, using a ChatCompletion model to select the next speaker after each message. + If an :class:`~autogen_agentchat.base.ChatAgent` is a participant, + the :class:`~autogen_agentchat.messages.BaseChatMessage` from the agent response's + :attr:`~autogen_agentchat.base.Response.chat_message` will be published + to other participants in the group chat. + + If a :class:`~autogen_agentchat.base.Team` is a participant, + the :class:`~autogen_agentchat.messages.BaseChatMessage` + from the team result' :attr:`~autogen_agentchat.base.TaskResult.messages` will be published + to other participants in the group chat. + Args: - participants (List[ChatAgent]): The participants in the group chat, + participants (List[ChatAgent | Team]): The participants in the group chat, must have unique names and at least two participants. model_client (ChatCompletionClient): The ChatCompletion model client used to select the next speaker. + name (str | None, optional): The name of the group chat, using + :attr:`~autogen_agentchat.teams.SelectorGroupChat.DEFAULT_NAME` if not provided. + The name is used by a parent team to identify this group chat so it must + be unique within the parent team. + description (str | None, optional): The description of the group chat, using + :attr:`~autogen_agentchat.teams.SelectorGroupChat.DEFAULT_DESCRIPTION` if not provided. termination_condition (TerminationCondition, optional): The termination condition for the group chat. Defaults to None. Without a termination condition, the group chat will run indefinitely. max_turns (int, optional): The maximum number of turns in the group chat before stopping. Defaults to None, meaning no limit. selector_prompt (str, optional): The prompt template to use for selecting the next speaker. Available fields: '{roles}', '{participants}', and '{history}'. + `{participants}` is the names of candidates for selection. The format is `["", "", ...]`. + `{roles}` is a newline-separated list of names and descriptions of the candidate agents. The format for each line is: `" : "`. + `{history}` is the conversation history formatted as a double newline separated of names and message content. The format for each message is: `" : "`. allow_repeated_speaker (bool, optional): Whether to include the previous speaker in the list of candidates to be selected for the next turn. Defaults to False. The model may still select the previous speaker -- a warning will be logged if this happens. max_selector_attempts (int, optional): The maximum number of attempts to select a speaker using the model. Defaults to 3. If the model fails to select a speaker after the maximum number of attempts, the previous speaker will be used if available, otherwise the first participant will be used. - selector_func (Callable[[Sequence[AgentEvent | ChatMessage]], str | None], Callable[[Sequence[AgentEvent | ChatMessage]], Awaitable[str | None]], optional): A custom selector + selector_func (Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], str | None], Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[str | None]], optional): A custom selector function that takes the conversation history and returns the name of the next speaker. If provided, this function will be used to override the model to select the next speaker. If the function returns None, the model will be used to select the next speaker. - candidate_func (Callable[[Sequence[AgentEvent | ChatMessage]], List[str]], Callable[[Sequence[AgentEvent | ChatMessage]], Awaitable[List[str]]], optional): + NOTE: `selector_func` is not serializable and will be ignored during serialization and deserialization process. + candidate_func (Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], List[str]], Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[List[str]]], optional): A custom function that takes the conversation history and returns a filtered list of candidates for the next speaker selection using model. If the function returns an empty list or `None`, `SelectorGroupChat` will raise a `ValueError`. This function is only used if `selector_func` is not set. The `allow_repeated_speaker` will be ignored if set. - + custom_message_types (List[type[BaseAgentEvent | BaseChatMessage]], optional): A list of custom message types that will be used in the group chat. + If you are using custom message types or your agents produces custom message types, you need to specify them here. + Make sure your custom message types are subclasses of :class:`~autogen_agentchat.messages.BaseAgentEvent` or :class:`~autogen_agentchat.messages.BaseChatMessage`. + emit_team_events (bool, optional): Whether to emit team events through :meth:`BaseGroupChat.run_stream`. Defaults to False. + model_client_streaming (bool, optional): Whether to use streaming for the model client. (This is useful for reasoning models like QwQ). Defaults to False. + model_context (ChatCompletionContext | None, optional): The model context for storing and retrieving + :class:`~autogen_core.models.LLMMessage`. It can be preloaded with initial messages. Messages stored in model context will be used for speaker selection. The initial messages will be cleared when the team is reset. Raises: ValueError: If the number of participants is less than two or if the selector prompt is invalid. @@ -387,7 +483,7 @@ async def book_trip() -> str: from autogen_agentchat.teams import SelectorGroupChat from autogen_agentchat.conditions import TextMentionTermination from autogen_agentchat.ui import Console - from autogen_agentchat.messages import AgentEvent, ChatMessage + from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage async def main() -> None: @@ -413,8 +509,8 @@ def check_calculation(x: int, y: int, answer: int) -> str: system_message="Check the answer and respond with 'Correct!' or 'Incorrect!'", ) - def selector_func(messages: Sequence[AgentEvent | ChatMessage]) -> str | None: - if len(messages) == 1 or messages[-1].content == "Incorrect!": + def selector_func(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> str | None: + if len(messages) == 1 or messages[-1].to_text() == "Incorrect!": return "Agent1" if messages[-1].source == "Agent1": return "Agent2" @@ -431,17 +527,80 @@ def selector_func(messages: Sequence[AgentEvent | ChatMessage]) -> str | None: await Console(team.run_stream(task="What is 1 + 1?")) + asyncio.run(main()) + + A team with custom model context: + + .. code-block:: python + + import asyncio + + from autogen_core.model_context import BufferedChatCompletionContext + from autogen_ext.models.openai import OpenAIChatCompletionClient + + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.conditions import TextMentionTermination + from autogen_agentchat.teams import SelectorGroupChat + from autogen_agentchat.ui import Console + + + async def main() -> None: + model_client = OpenAIChatCompletionClient(model="gpt-4o") + model_context = BufferedChatCompletionContext(buffer_size=5) + + async def lookup_hotel(location: str) -> str: + return f"Here are some hotels in {location}: hotel1, hotel2, hotel3." + + async def lookup_flight(origin: str, destination: str) -> str: + return f"Here are some flights from {origin} to {destination}: flight1, flight2, flight3." + + async def book_trip() -> str: + return "Your trip is booked!" + + travel_advisor = AssistantAgent( + "Travel_Advisor", + model_client, + tools=[book_trip], + description="Helps with travel planning.", + ) + hotel_agent = AssistantAgent( + "Hotel_Agent", + model_client, + tools=[lookup_hotel], + description="Helps with hotel booking.", + ) + flight_agent = AssistantAgent( + "Flight_Agent", + model_client, + tools=[lookup_flight], + description="Helps with flight booking.", + ) + termination = TextMentionTermination("TERMINATE") + team = SelectorGroupChat( + [travel_advisor, hotel_agent, flight_agent], + model_client=model_client, + termination_condition=termination, + model_context=model_context, + ) + await Console(team.run_stream(task="Book a 3-day trip to new york.")) + + asyncio.run(main()) """ component_config_schema = SelectorGroupChatConfig component_provider_override = "autogen_agentchat.teams.SelectorGroupChat" + DEFAULT_NAME = "SelectorGroupChat" + DEFAULT_DESCRIPTION = "A team of agents." + def __init__( self, - participants: List[ChatAgent], + participants: List[ChatAgent | Team], model_client: ChatCompletionClient, *, + name: str | None = None, + description: str | None = None, termination_condition: TerminationCondition | None = None, max_turns: int | None = None, runtime: AgentRuntime | None = None, @@ -457,14 +616,22 @@ def __init__( max_selector_attempts: int = 3, selector_func: Optional[SelectorFuncType] = None, candidate_func: Optional[CandidateFuncType] = None, + custom_message_types: List[type[BaseAgentEvent | BaseChatMessage]] | None = None, + emit_team_events: bool = False, + model_client_streaming: bool = False, + model_context: ChatCompletionContext | None = None, ): super().__init__( - participants, + name=name or self.DEFAULT_NAME, + description=description or self.DEFAULT_DESCRIPTION, + participants=participants, group_chat_manager_name="SelectorGroupChatManager", group_chat_manager_class=SelectorGroupChatManager, termination_condition=termination_condition, max_turns=max_turns, runtime=runtime, + custom_message_types=custom_message_types, + emit_team_events=emit_team_events, ) # Validate the participants. if len(participants) < 2: @@ -475,6 +642,8 @@ def __init__( self._selector_func = selector_func self._max_selector_attempts = max_selector_attempts self._candidate_func = candidate_func + self._model_client_streaming = model_client_streaming + self._model_context = model_context def _create_group_chat_manager_factory( self, @@ -484,9 +653,10 @@ def _create_group_chat_manager_factory( participant_topic_types: List[str], participant_names: List[str], participant_descriptions: List[str], - output_message_queue: asyncio.Queue[AgentEvent | ChatMessage | GroupChatTermination], + output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination], termination_condition: TerminationCondition | None, max_turns: int | None, + message_factory: MessageFactory, ) -> Callable[[], BaseGroupChatManager]: return lambda: SelectorGroupChatManager( name, @@ -498,16 +668,22 @@ def _create_group_chat_manager_factory( output_message_queue, termination_condition, max_turns, + message_factory, self._model_client, self._selector_prompt, self._allow_repeated_speaker, self._selector_func, self._max_selector_attempts, self._candidate_func, + self._emit_team_events, + self._model_context, + self._model_client_streaming, ) def _to_config(self) -> SelectorGroupChatConfig: return SelectorGroupChatConfig( + name=self._name, + description=self._description, participants=[participant.dump_component() for participant in self._participants], model_client=self._model_client.dump_component(), termination_condition=self._termination_condition.dump_component() if self._termination_condition else None, @@ -516,13 +692,28 @@ def _to_config(self) -> SelectorGroupChatConfig: allow_repeated_speaker=self._allow_repeated_speaker, max_selector_attempts=self._max_selector_attempts, # selector_func=self._selector_func.dump_component() if self._selector_func else None, + emit_team_events=self._emit_team_events, + model_client_streaming=self._model_client_streaming, + model_context=self._model_context.dump_component() if self._model_context else None, ) @classmethod def _from_config(cls, config: SelectorGroupChatConfig) -> Self: + participants: List[ChatAgent | Team] = [] + for participant in config.participants: + if participant.component_type == ChatAgent.component_type: + participants.append(ChatAgent.load_component(participant)) + elif participant.component_type == Team.component_type: + participants.append(Team.load_component(participant)) + else: + raise ValueError( + f"Invalid participant component type: {participant.component_type}. " "Expected ChatAgent or Team." + ) return cls( - participants=[BaseChatAgent.load_component(participant) for participant in config.participants], + participants=participants, model_client=ChatCompletionClient.load_component(config.model_client), + name=config.name, + description=config.description, termination_condition=TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None, @@ -530,7 +721,10 @@ def _from_config(cls, config: SelectorGroupChatConfig) -> Self: selector_prompt=config.selector_prompt, allow_repeated_speaker=config.allow_repeated_speaker, max_selector_attempts=config.max_selector_attempts, - # selector_func=ComponentLoader.load_component(config.selector_func, Callable[[Sequence[AgentEvent | ChatMessage]], str | None]) + # selector_func=ComponentLoader.load_component(config.selector_func, Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], str | None]) # if config.selector_func # else None, + emit_team_events=config.emit_team_events, + model_client_streaming=config.model_client_streaming, + model_context=ChatCompletionContext.load_component(config.model_context) if config.model_context else None, ) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py index f28965712b71..c9b495083939 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py @@ -1,11 +1,11 @@ import asyncio -from typing import Any, Callable, List, Mapping +from typing import Any, Callable, List, Mapping, Sequence from autogen_core import AgentRuntime, Component, ComponentModel from pydantic import BaseModel from ...base import ChatAgent, TerminationCondition -from ...messages import AgentEvent, ChatMessage, HandoffMessage +from ...messages import BaseAgentEvent, BaseChatMessage, HandoffMessage, MessageFactory from ...state import SwarmManagerState from ._base_group_chat import BaseGroupChat from ._base_group_chat_manager import BaseGroupChatManager @@ -23,9 +23,11 @@ def __init__( participant_topic_types: List[str], participant_names: List[str], participant_descriptions: List[str], - output_message_queue: asyncio.Queue[AgentEvent | ChatMessage | GroupChatTermination], + output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination], termination_condition: TerminationCondition | None, max_turns: int | None, + message_factory: MessageFactory, + emit_team_events: bool, ) -> None: super().__init__( name, @@ -37,10 +39,12 @@ def __init__( output_message_queue, termination_condition, max_turns, + message_factory, + emit_team_events, ) self._current_speaker = self._participant_names[0] - async def validate_group_state(self, messages: List[ChatMessage] | None) -> None: + async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None: """Validate the start messages for the group chat.""" # Check if any of the start messages is a handoff message. if messages: @@ -75,22 +79,27 @@ async def reset(self) -> None: await self._termination_condition.reset() self._current_speaker = self._participant_names[0] - async def select_speaker(self, thread: List[AgentEvent | ChatMessage]) -> str: + async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str: """Select a speaker from the participants based on handoff message. - Looks for the last handoff message in the thread to determine the next speaker.""" + Looks for the last handoff message in the thread to determine the next speaker. + + .. note:: + + This method always returns a single speaker. + """ if len(thread) == 0: - return self._current_speaker + return [self._current_speaker] for message in reversed(thread): if isinstance(message, HandoffMessage): self._current_speaker = message.target # The latest handoff message should always target a valid participant. assert self._current_speaker in self._participant_names - return self._current_speaker + return [self._current_speaker] return self._current_speaker async def save_state(self) -> Mapping[str, Any]: state = SwarmManagerState( - message_thread=list(self._message_thread), + message_thread=[msg.dump() for msg in self._message_thread], current_turn=self._current_turn, current_speaker=self._current_speaker, ) @@ -98,7 +107,7 @@ async def save_state(self) -> Mapping[str, Any]: async def load_state(self, state: Mapping[str, Any]) -> None: swarm_state = SwarmManagerState.model_validate(state) - self._message_thread = list(swarm_state.message_thread) + self._message_thread = [self._message_factory.create(message) for message in swarm_state.message_thread] self._current_turn = swarm_state.current_turn self._current_speaker = swarm_state.current_speaker @@ -106,9 +115,12 @@ async def load_state(self, state: Mapping[str, Any]) -> None: class SwarmConfig(BaseModel): """The declarative configuration for Swarm.""" + name: str | None = None + description: str | None = None participants: List[ComponentModel] termination_condition: ComponentModel | None = None max_turns: int | None = None + emit_team_events: bool = False class Swarm(BaseGroupChat, Component[SwarmConfig]): @@ -119,11 +131,24 @@ class Swarm(BaseGroupChat, Component[SwarmConfig]): sent by the current speaker. If no handoff message is sent, the current speaker continues to be the speaker. + .. note:: + + Unlike :class:`~autogen_agentchat.teams.RoundRobinGroupChat` and + :class:`~autogen_agentchat.teams.SelectorGroupChat`, this group chat + team does not support inner teams as participants. + Args: participants (List[ChatAgent]): The agents participating in the group chat. The first agent in the list is the initial speaker. + name (str | None, optional): The name of the group chat, using :attr:`~autogen_agentchat.teams.Swarm.DEFAULT_NAME` if not provided. + The name is used by a parent team to identify this group chat so it must be unique within the parent team. + description (str | None, optional): The description of the group chat, using :attr:`~autogen_agentchat.teams.Swarm.DEFAULT_DESCRIPTION` if not provided. termination_condition (TerminationCondition, optional): The termination condition for the group chat. Defaults to None. Without a termination condition, the group chat will run indefinitely. max_turns (int, optional): The maximum number of turns in the group chat before stopping. Defaults to None, meaning no limit. + custom_message_types (List[type[BaseAgentEvent | BaseChatMessage]], optional): A list of custom message types that will be used in the group chat. + If you are using custom message types or your agents produces custom message types, you need to specify them here. + Make sure your custom message types are subclasses of :class:`~autogen_agentchat.messages.BaseAgentEvent` or :class:`~autogen_agentchat.messages.BaseChatMessage`. + emit_team_events (bool, optional): Whether to emit team events through :meth:`BaseGroupChat.run_stream`. Defaults to False. Basic example: @@ -202,25 +227,39 @@ async def main() -> None: component_config_schema = SwarmConfig component_provider_override = "autogen_agentchat.teams.Swarm" - # TODO: Add * to the constructor to separate the positional parameters from the kwargs. - # This may be a breaking change so let's wait until a good time to do it. + DEFAULT_NAME = "Swarm" + DEFAULT_DESCRIPTION = "A team of agents." + def __init__( self, participants: List[ChatAgent], + *, + name: str | None = None, + description: str | None = None, termination_condition: TerminationCondition | None = None, max_turns: int | None = None, runtime: AgentRuntime | None = None, + custom_message_types: List[type[BaseAgentEvent | BaseChatMessage]] | None = None, + emit_team_events: bool = False, ) -> None: + for participant in participants: + if not isinstance(participant, ChatAgent): + raise TypeError(f"Participant {participant} must be a ChatAgent.") super().__init__( - participants, + name=name or self.DEFAULT_NAME, + description=description or self.DEFAULT_DESCRIPTION, + participants=[participant for participant in participants], group_chat_manager_name="SwarmGroupChatManager", group_chat_manager_class=SwarmGroupChatManager, termination_condition=termination_condition, max_turns=max_turns, runtime=runtime, + custom_message_types=custom_message_types, + emit_team_events=emit_team_events, ) # The first participant must be able to produce handoff messages. first_participant = self._participants[0] + assert isinstance(first_participant, ChatAgent) if HandoffMessage not in first_participant.produced_message_types: raise ValueError("The first participant must be able to produce a handoff messages.") @@ -232,9 +271,10 @@ def _create_group_chat_manager_factory( participant_topic_types: List[str], participant_names: List[str], participant_descriptions: List[str], - output_message_queue: asyncio.Queue[AgentEvent | ChatMessage | GroupChatTermination], + output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination], termination_condition: TerminationCondition | None, max_turns: int | None, + message_factory: MessageFactory, ) -> Callable[[], SwarmGroupChatManager]: def _factory() -> SwarmGroupChatManager: return SwarmGroupChatManager( @@ -247,6 +287,8 @@ def _factory() -> SwarmGroupChatManager: output_message_queue, termination_condition, max_turns, + message_factory, + self._emit_team_events, ) return _factory @@ -255,9 +297,12 @@ def _to_config(self) -> SwarmConfig: participants = [participant.dump_component() for participant in self._participants] termination_condition = self._termination_condition.dump_component() if self._termination_condition else None return SwarmConfig( + name=self._name, + description=self._description, participants=participants, termination_condition=termination_condition, max_turns=self._max_turns, + emit_team_events=self._emit_team_events, ) @classmethod @@ -266,4 +311,11 @@ def _from_config(cls, config: SwarmConfig) -> "Swarm": termination_condition = ( TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None ) - return cls(participants, termination_condition=termination_condition, max_turns=config.max_turns) + return cls( + participants, + name=config.name, + description=config.description, + termination_condition=termination_condition, + max_turns=config.max_turns, + emit_team_events=config.emit_team_events, + ) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/tools/__init__.py b/python/packages/autogen-agentchat/src/autogen_agentchat/tools/__init__.py new file mode 100644 index 000000000000..9884ddcd889d --- /dev/null +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/tools/__init__.py @@ -0,0 +1,4 @@ +from ._agent import AgentTool +from ._team import TeamTool + +__all__ = ["AgentTool", "TeamTool"] diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/tools/_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/tools/_agent.py new file mode 100644 index 000000000000..ba83bea6b61d --- /dev/null +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/tools/_agent.py @@ -0,0 +1,93 @@ +from autogen_core import Component, ComponentModel +from pydantic import BaseModel +from typing_extensions import Self + +from autogen_agentchat.agents import BaseChatAgent + +from ._task_runner_tool import TaskRunnerTool + + +class AgentToolConfig(BaseModel): + """Configuration for the AgentTool.""" + + agent: ComponentModel + """The agent to be used for running the task.""" + + return_value_as_last_message: bool = False + """Whether to return the value as the last message of the task result.""" + + +class AgentTool(TaskRunnerTool, Component[AgentToolConfig]): + """Tool that can be used to run a task using an agent. + + The tool returns the result of the task execution as a :class:`~autogen_agentchat.base.TaskResult` object. + + .. important:: + When using AgentTool, you **must** disable parallel tool calls in the model client configuration + to avoid concurrency issues. Agents cannot run concurrently as they maintain internal state + that would conflict with parallel execution. For example, set ``parallel_tool_calls=False`` + for :class:`~autogen_ext.models.openai.OpenAIChatCompletionClient` and + :class:`~autogen_ext.models.openai.AzureOpenAIChatCompletionClient`. + + Args: + agent (BaseChatAgent): The agent to be used for running the task. + return_value_as_last_message (bool): Whether to use the last message content of the task result + as the return value of the tool in :meth:`~autogen_agentchat.tools.TaskRunnerTool.return_value_as_string`. + If set to True, the last message content will be returned as a string. + If set to False, the tool will return all messages in the task result as a string concatenated together, + with each message prefixed by its source (e.g., "writer: ...", "assistant: ..."). + + Example: + + .. code-block:: python + + import asyncio + + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.tools import AgentTool + from autogen_agentchat.ui import Console + from autogen_ext.models.openai import OpenAIChatCompletionClient + + + async def main() -> None: + model_client = OpenAIChatCompletionClient(model="gpt-4.1") + writer = AssistantAgent( + name="writer", + description="A writer agent for generating text.", + model_client=model_client, + system_message="Write well.", + ) + writer_tool = AgentTool(agent=writer) + + # Create model client with parallel tool calls disabled for the main agent + main_model_client = OpenAIChatCompletionClient(model="gpt-4.1", parallel_tool_calls=False) + assistant = AssistantAgent( + name="assistant", + model_client=main_model_client, + tools=[writer_tool], + system_message="You are a helpful assistant.", + ) + await Console(assistant.run_stream(task="Write a poem about the sea.")) + + + asyncio.run(main()) + """ + + component_config_schema = AgentToolConfig + component_provider_override = "autogen_agentchat.tools.AgentTool" + + def __init__(self, agent: BaseChatAgent, return_value_as_last_message: bool = False) -> None: + self._agent = agent + super().__init__( + agent, agent.name, agent.description, return_value_as_last_message=return_value_as_last_message + ) + + def _to_config(self) -> AgentToolConfig: + return AgentToolConfig( + agent=self._agent.dump_component(), + return_value_as_last_message=self._return_value_as_last_message, + ) + + @classmethod + def _from_config(cls, config: AgentToolConfig) -> Self: + return cls(BaseChatAgent.load_component(config.agent), config.return_value_as_last_message) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/tools/_task_runner_tool.py b/python/packages/autogen-agentchat/src/autogen_agentchat/tools/_task_runner_tool.py new file mode 100644 index 000000000000..0db95ed2b2a0 --- /dev/null +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/tools/_task_runner_tool.py @@ -0,0 +1,71 @@ +from abc import ABC +from typing import Annotated, Any, AsyncGenerator, List, Mapping + +from autogen_core import CancellationToken +from autogen_core.tools import BaseStreamTool +from pydantic import BaseModel + +from ..agents import BaseChatAgent +from ..base import TaskResult +from ..messages import BaseAgentEvent, BaseChatMessage +from ..teams import BaseGroupChat + + +class TaskRunnerToolArgs(BaseModel): + """Input for the TaskRunnerTool.""" + + task: Annotated[str, "The task to be executed."] + + +class TaskRunnerTool(BaseStreamTool[TaskRunnerToolArgs, BaseAgentEvent | BaseChatMessage, TaskResult], ABC): + """An base class for tool that can be used to run a task using a team or an agent.""" + + component_type = "tool" + + def __init__( + self, + task_runner: BaseGroupChat | BaseChatAgent, + name: str, + description: str, + return_value_as_last_message: bool, + ) -> None: + self._task_runner = task_runner + self._return_value_as_last_message = return_value_as_last_message + super().__init__( + args_type=TaskRunnerToolArgs, + return_type=TaskResult, + name=name, + description=description, + ) + + async def run(self, args: TaskRunnerToolArgs, cancellation_token: CancellationToken) -> TaskResult: + """Run the task and return the result.""" + return await self._task_runner.run(task=args.task, cancellation_token=cancellation_token) + + async def run_stream( + self, args: TaskRunnerToolArgs, cancellation_token: CancellationToken + ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None]: + """Run the task and yield events or messages as they are produced, the final :class:`TaskResult` + will be yielded at the end.""" + async for event in self._task_runner.run_stream(task=args.task, cancellation_token=cancellation_token): + yield event + + def return_value_as_string(self, value: TaskResult) -> str: + """Convert the task result to a string.""" + if self._return_value_as_last_message: + if value.messages and isinstance(value.messages[-1], BaseChatMessage): + return value.messages[-1].to_model_text() + raise ValueError("The last message is not a BaseChatMessage.") + parts: List[str] = [] + for message in value.messages: + if isinstance(message, BaseChatMessage): + if message.source == "user": + continue + parts.append(f"{message.source}: {message.to_model_text()}") + return "\n\n".join(parts) + + async def save_state_json(self) -> Mapping[str, Any]: + return await self._task_runner.save_state() + + async def load_state_json(self, state: Mapping[str, Any]) -> None: + await self._task_runner.load_state(state) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/tools/_team.py b/python/packages/autogen-agentchat/src/autogen_agentchat/tools/_team.py new file mode 100644 index 000000000000..9c8ecf1b0634 --- /dev/null +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/tools/_team.py @@ -0,0 +1,133 @@ +from autogen_core import Component, ComponentModel +from pydantic import BaseModel +from typing_extensions import Self + +from autogen_agentchat.teams import BaseGroupChat + +from ._task_runner_tool import TaskRunnerTool + + +class TeamToolConfig(BaseModel): + """Configuration for the TeamTool.""" + + name: str + """The name of the tool.""" + description: str + """The name and description of the tool.""" + team: ComponentModel + """The team to be used for running the task.""" + return_value_as_last_message: bool = False + """Whether to return the value as the last message of the task result.""" + + +class TeamTool(TaskRunnerTool, Component[TeamToolConfig]): + """Tool that can be used to run a task. + + The tool returns the result of the task execution as a :class:`~autogen_agentchat.base.TaskResult` object. + + .. important:: + When using TeamTool, you **must** disable parallel tool calls in the model client configuration + to avoid concurrency issues. Teams cannot run concurrently as they maintain internal state + that would conflict with parallel execution. For example, set ``parallel_tool_calls=False`` + for :class:`~autogen_ext.models.openai.OpenAIChatCompletionClient` and + :class:`~autogen_ext.models.openai.AzureOpenAIChatCompletionClient`. + + Args: + team (BaseGroupChat): The team to be used for running the task. + name (str): The name of the tool. + description (str): The description of the tool. + return_value_as_last_message (bool): Whether to use the last message content of the task result + as the return value of the tool in :meth:`~autogen_agentchat.tools.TaskRunnerTool.return_value_as_string`. + If set to True, the last message content will be returned as a string. + If set to False, the tool will return all messages in the task result as a string concatenated together, + with each message prefixed by its source (e.g., "writer: ...", "assistant: ..."). + + Example: + + .. code-block:: python + + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.conditions import SourceMatchTermination + from autogen_agentchat.teams import RoundRobinGroupChat + from autogen_agentchat.tools import TeamTool + from autogen_agentchat.ui import Console + from autogen_ext.models.openai import OpenAIChatCompletionClient + + + async def main() -> None: + # Disable parallel tool calls when using TeamTool + model_client = OpenAIChatCompletionClient(model="gpt-4.1") + + writer = AssistantAgent(name="writer", model_client=model_client, system_message="You are a helpful assistant.") + reviewer = AssistantAgent( + name="reviewer", model_client=model_client, system_message="You are a critical reviewer." + ) + summarizer = AssistantAgent( + name="summarizer", + model_client=model_client, + system_message="You combine the review and produce a revised response.", + ) + team = RoundRobinGroupChat( + [writer, reviewer, summarizer], termination_condition=SourceMatchTermination(sources=["summarizer"]) + ) + + # Create a TeamTool that uses the team to run tasks, returning the last message as the result. + tool = TeamTool( + team=team, + name="writing_team", + description="A tool for writing tasks.", + return_value_as_last_message=True, + ) + + # Create model client with parallel tool calls disabled for the main agent + main_model_client = OpenAIChatCompletionClient(model="gpt-4.1", parallel_tool_calls=False) + main_agent = AssistantAgent( + name="main_agent", + model_client=main_model_client, + system_message="You are a helpful assistant that can use the writing tool.", + tools=[tool], + ) + # For handling each events manually. + # async for message in main_agent.run_stream( + # task="Write a short story about a robot learning to love.", + # ): + # print(message) + # Use Console to display the messages in a more readable format. + await Console( + main_agent.run_stream( + task="Write a short story about a robot learning to love.", + ) + ) + + + if __name__ == "__main__": + import asyncio + + asyncio.run(main()) + """ + + component_config_schema = TeamToolConfig + component_provider_override = "autogen_agentchat.tools.TeamTool" + + def __init__( + self, team: BaseGroupChat, name: str, description: str, return_value_as_last_message: bool = False + ) -> None: + self._team = team + super().__init__(team, name, description, return_value_as_last_message=return_value_as_last_message) + + def _to_config(self) -> TeamToolConfig: + return TeamToolConfig( + name=self._name, + description=self._description, + team=self._team.dump_component(), + return_value_as_last_message=self._return_value_as_last_message, + ) + + @classmethod + def _from_config(cls, config: TeamToolConfig) -> Self: + return cls( + BaseGroupChat.load_component(config.team), + config.name, + config.description, + config.return_value_as_last_message, + ) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/ui/_console.py b/python/packages/autogen-agentchat/src/autogen_agentchat/ui/_console.py index 0a95c842ea08..12fba2b489c1 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/ui/_console.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/ui/_console.py @@ -5,14 +5,14 @@ from inspect import iscoroutinefunction from typing import AsyncGenerator, Awaitable, Callable, Dict, List, Optional, TypeVar, Union, cast -from autogen_core import CancellationToken, Image +from autogen_core import CancellationToken from autogen_core.models import RequestUsage from autogen_agentchat.agents import UserProxyAgent from autogen_agentchat.base import Response, TaskResult from autogen_agentchat.messages import ( - AgentEvent, - ChatMessage, + BaseAgentEvent, + BaseChatMessage, ModelClientStreamingChunkEvent, MultiModalMessage, UserInputRequestedEvent, @@ -80,7 +80,7 @@ def aprint(output: str, end: str = "\n", flush: bool = False) -> Awaitable[None] async def Console( - stream: AsyncGenerator[AgentEvent | ChatMessage | T, None], + stream: AsyncGenerator[BaseAgentEvent | BaseChatMessage | T, None], *, no_inline_images: bool = False, output_stats: bool = False, @@ -97,7 +97,7 @@ async def Console( It will be improved in future releases. Args: - stream (AsyncGenerator[AgentEvent | ChatMessage | TaskResult, None] | AsyncGenerator[AgentEvent | ChatMessage | Response, None]): Message stream to render. + stream (AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None] | AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]): Message stream to render. This can be from :meth:`~autogen_agentchat.base.TaskRunner.run_stream` or :meth:`~autogen_agentchat.base.ChatAgent.on_messages_stream`. no_inline_images (bool, optional): If terminal is iTerm2 will render images inline. Use this to disable this behavior. Defaults to False. output_stats (bool, optional): (Experimental) If True, will output a summary of the messages and inline token usage info. Defaults to False. @@ -135,7 +135,11 @@ async def Console( duration = time.time() - start_time # Print final response. - output = f"{'-' * 10} {message.chat_message.source} {'-' * 10}\n{_message_to_str(message.chat_message, render_image_iterm=render_image_iterm)}\n" + if isinstance(message.chat_message, MultiModalMessage): + final_content = message.chat_message.to_text(iterm=render_image_iterm) + else: + final_content = message.chat_message.to_text() + output = f"{'-' * 10} {message.chat_message.source} {'-' * 10}\n{final_content}\n" if message.chat_message.models_usage: if output_stats: output += f"[Prompt tokens: {message.chat_message.models_usage.prompt_tokens}, Completion tokens: {message.chat_message.models_usage.completion_tokens}]\n" @@ -166,21 +170,24 @@ async def Console( user_input_manager.notify_event_received(message.request_id) else: # Cast required for mypy to be happy - message = cast(AgentEvent | ChatMessage, message) # type: ignore + message = cast(BaseAgentEvent | BaseChatMessage, message) # type: ignore if not streaming_chunks: # Print message sender. - await aprint(f"{'-' * 10} {message.source} {'-' * 10}", end="\n", flush=True) + await aprint( + f"{'-' * 10} {message.__class__.__name__} ({message.source}) {'-' * 10}", end="\n", flush=True + ) if isinstance(message, ModelClientStreamingChunkEvent): - await aprint(message.content, end="") + await aprint(message.to_text(), end="", flush=True) streaming_chunks.append(message.content) else: if streaming_chunks: streaming_chunks.clear() # Chunked messages are already printed, so we just print a newline. await aprint("", end="\n", flush=True) + elif isinstance(message, MultiModalMessage): + await aprint(message.to_text(iterm=render_image_iterm), end="\n", flush=True) else: - # Print message content. - await aprint(_message_to_str(message, render_image_iterm=render_image_iterm), end="\n", flush=True) + await aprint(message.to_text(), end="\n", flush=True) if message.models_usage: if output_stats: await aprint( @@ -195,25 +202,3 @@ async def Console( raise ValueError("No TaskResult or Response was processed.") return last_processed - - -# iTerm2 image rendering protocol: https://iterm2.com/documentation-images.html -def _image_to_iterm(image: Image) -> str: - image_data = image.to_base64() - return f"\033]1337;File=inline=1:{image_data}\a\n" - - -def _message_to_str(message: AgentEvent | ChatMessage, *, render_image_iterm: bool = False) -> str: - if isinstance(message, MultiModalMessage): - result: List[str] = [] - for c in message.content: - if isinstance(c, str): - result.append(c) - else: - if render_image_iterm: - result.append(_image_to_iterm(c)) - else: - result.append("") - return "\n".join(result) - else: - return f"{message.content}" diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/utils/_utils.py b/python/packages/autogen-agentchat/src/autogen_agentchat/utils/_utils.py index 6de1178645fc..738b72e9b329 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/utils/_utils.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/utils/_utils.py @@ -2,18 +2,24 @@ from autogen_core import FunctionCall, Image from autogen_core.models import FunctionExecutionResult, LLMMessage, UserMessage +from pydantic import BaseModel # Type aliases for convenience +_StructuredContent = BaseModel _UserContent = Union[str, List[Union[str, Image]]] _AssistantContent = Union[str, List[FunctionCall]] _FunctionExecutionContent = List[FunctionExecutionResult] _SystemContent = str -def content_to_str(content: _UserContent | _AssistantContent | _FunctionExecutionContent | _SystemContent) -> str: +def content_to_str( + content: _UserContent | _AssistantContent | _FunctionExecutionContent | _SystemContent | _StructuredContent, +) -> str: """Convert the content of an LLMMessage to a string.""" if isinstance(content, str): return content + elif isinstance(content, BaseModel): + return content.model_dump_json() else: result: List[str] = [] for c in content: diff --git a/python/packages/autogen-agentchat/tests/test_agent.py b/python/packages/autogen-agentchat/tests/test_agent.py new file mode 100644 index 000000000000..605f85448aa6 --- /dev/null +++ b/python/packages/autogen-agentchat/tests/test_agent.py @@ -0,0 +1,126 @@ +import pytest +from autogen_agentchat.agents import ( + AssistantAgent, + CodeExecutorAgent, + SocietyOfMindAgent, +) +from autogen_agentchat.teams import RoundRobinGroupChat +from autogen_core.model_context import ( + BufferedChatCompletionContext, + ChatCompletionContext, + HeadAndTailChatCompletionContext, + TokenLimitedChatCompletionContext, + UnboundedChatCompletionContext, +) +from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor +from autogen_ext.models.replay import ReplayChatCompletionClient + + +@pytest.mark.parametrize( + "model_context_class", + [ + UnboundedChatCompletionContext(), + BufferedChatCompletionContext(buffer_size=5), + TokenLimitedChatCompletionContext(model_client=ReplayChatCompletionClient([]), token_limit=5), + HeadAndTailChatCompletionContext(head_size=3, tail_size=2), + ], +) +def test_serialize_and_deserialize_model_context_on_assistant_agent(model_context_class: ChatCompletionContext) -> None: + """Test the serialization and deserialization of the message context on the AssistantAgent.""" + agent = AssistantAgent( + name="assistant", + model_client=ReplayChatCompletionClient([]), + description="An assistant agent.", + model_context=model_context_class, + ) + + # Serialize the agent + serialized_agent = agent.dump_component() + # Deserialize the agent + deserialized_agent = AssistantAgent.load_component(serialized_agent) + + # Check that the deserialized agent has the same model context as the original agent + original_model_context = agent.model_context + deserialized_model_context = deserialized_agent.model_context + + assert isinstance(original_model_context, type(deserialized_model_context)) + assert isinstance(deserialized_model_context, type(original_model_context)) + assert original_model_context.dump_component() == deserialized_model_context.dump_component() + + +@pytest.mark.parametrize( + "model_context_class", + [ + UnboundedChatCompletionContext(), + BufferedChatCompletionContext(buffer_size=5), + TokenLimitedChatCompletionContext(model_client=ReplayChatCompletionClient([]), token_limit=5), + HeadAndTailChatCompletionContext(head_size=3, tail_size=2), + ], +) +def test_serialize_and_deserialize_model_context_on_society_of_mind_agent( + model_context_class: ChatCompletionContext, +) -> None: + """Test the serialization and deserialization of the message context on the AssistantAgent.""" + agent1 = AssistantAgent( + name="assistant1", model_client=ReplayChatCompletionClient([]), description="An assistant agent." + ) + agent2 = AssistantAgent( + name="assistant2", model_client=ReplayChatCompletionClient([]), description="An assistant agent." + ) + team = RoundRobinGroupChat( + participants=[agent1, agent2], + ) + agent = SocietyOfMindAgent( + name="assistant", + model_client=ReplayChatCompletionClient([]), + description="An assistant agent.", + team=team, + model_context=model_context_class, + ) + + # Serialize the agent + serialized_agent = agent.dump_component() + # Deserialize the agent + deserialized_agent = SocietyOfMindAgent.load_component(serialized_agent) + + # Check that the deserialized agent has the same model context as the original agent + original_model_context = agent.model_context + deserialized_model_context = deserialized_agent.model_context + + assert isinstance(original_model_context, type(deserialized_model_context)) + assert isinstance(deserialized_model_context, type(original_model_context)) + assert original_model_context.dump_component() == deserialized_model_context.dump_component() + + +@pytest.mark.parametrize( + "model_context_class", + [ + UnboundedChatCompletionContext(), + BufferedChatCompletionContext(buffer_size=5), + TokenLimitedChatCompletionContext(model_client=ReplayChatCompletionClient([]), token_limit=5), + HeadAndTailChatCompletionContext(head_size=3, tail_size=2), + ], +) +def test_serialize_and_deserialize_model_context_on_code_executor_agent( + model_context_class: ChatCompletionContext, +) -> None: + """Test the serialization and deserialization of the message context on the AssistantAgent.""" + agent = CodeExecutorAgent( + name="assistant", + code_executor=LocalCommandLineCodeExecutor(), + description="An assistant agent.", + model_context=model_context_class, + ) + + # Serialize the agent + serialized_agent = agent.dump_component() + # Deserialize the agent + deserialized_agent = CodeExecutorAgent.load_component(serialized_agent) + + # Check that the deserialized agent has the same model context as the original agent + original_model_context = agent.model_context + deserialized_model_context = deserialized_agent.model_context + + assert isinstance(original_model_context, type(deserialized_model_context)) + assert isinstance(deserialized_model_context, type(original_model_context)) + assert original_model_context.dump_component() == deserialized_model_context.dump_component() diff --git a/python/packages/autogen-agentchat/tests/test_assistant_agent.py b/python/packages/autogen-agentchat/tests/test_assistant_agent.py index db4fe42b73c5..935f2471045c 100644 --- a/python/packages/autogen-agentchat/tests/test_assistant_agent.py +++ b/python/packages/autogen-agentchat/tests/test_assistant_agent.py @@ -1,943 +1,3562 @@ +"""Comprehensive tests for AssistantAgent functionality.""" + +# Standard library imports +import asyncio import json -import logging -from typing import Dict, List +import os +from typing import Any, List, Optional, Union, cast +from unittest.mock import AsyncMock, MagicMock, patch +# Third-party imports import pytest -from autogen_agentchat import EVENT_LOGGER_NAME + +# First-party imports from autogen_agentchat.agents import AssistantAgent -from autogen_agentchat.base import Handoff, TaskResult +from autogen_agentchat.agents._assistant_agent import AssistantAgentConfig +from autogen_agentchat.base import Handoff, Response, TaskResult from autogen_agentchat.messages import ( - ChatMessage, + BaseAgentEvent, + BaseChatMessage, HandoffMessage, MemoryQueryEvent, ModelClientStreamingChunkEvent, - MultiModalMessage, + StructuredMessage, TextMessage, ThoughtEvent, ToolCallExecutionEvent, ToolCallRequestEvent, ToolCallSummaryMessage, ) -from autogen_core import ComponentModel, FunctionCall, Image -from autogen_core.memory import ListMemory, Memory, MemoryContent, MemoryMimeType, MemoryQueryResult +from autogen_core import CancellationToken, ComponentModel, FunctionCall +from autogen_core.memory import Memory, MemoryContent, UpdateContextResult +from autogen_core.memory import MemoryQueryResult as MemoryQueryResultSet from autogen_core.model_context import BufferedChatCompletionContext from autogen_core.models import ( AssistantMessage, CreateResult, FunctionExecutionResult, - FunctionExecutionResultMessage, - LLMMessage, + ModelFamily, RequestUsage, SystemMessage, UserMessage, ) -from autogen_core.models._model_client import ModelFamily -from autogen_core.tools import BaseTool, FunctionTool +from autogen_ext.models.anthropic import AnthropicChatCompletionClient from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.models.replay import ReplayChatCompletionClient -from pydantic import BaseModel -from utils import FileLogHandler +from autogen_ext.tools.mcp import McpWorkbench, SseServerParams +from pydantic import BaseModel, ValidationError + + +def mock_tool_function(param: str) -> str: + """Mock tool function for testing. + + Args: + param: Input parameter to process + + Returns: + Formatted string with the input parameter + """ + return f"Tool executed with: {param}" + + +async def async_mock_tool_function(param: str) -> str: + """Async mock tool function for testing. -logger = logging.getLogger(EVENT_LOGGER_NAME) -logger.setLevel(logging.DEBUG) -logger.addHandler(FileLogHandler("test_assistant_agent.log")) + Args: + param: Input parameter to process + + Returns: + Formatted string with the input parameter + """ + return f"Async tool executed with: {param}" def _pass_function(input: str) -> str: + """Pass through function for testing. + + Args: + input: Input to pass through + + Returns: + The string "pass" + """ return "pass" -async def _fail_function(input: str) -> str: - return "fail" +def _echo_function(input: str) -> str: + """Echo function for testing. + Args: + input: Input to echo -async def _echo_function(input: str) -> str: + Returns: + The input string + """ return input +class MockMemory(Memory): + """Mock memory implementation for testing. + + A simple memory implementation that stores strings and provides basic memory operations + for testing purposes. + + Args: + contents: Optional list of initial memory contents + """ + + def __init__(self, contents: Optional[List[str]] = None) -> None: + """Initialize mock memory. + + Args: + contents: Optional list of initial memory contents + """ + self._contents: List[str] = contents or [] + + async def add(self, content: MemoryContent, cancellation_token: Optional[CancellationToken] = None) -> None: + """Add content to memory. + + Args: + content: Content to add to memory + cancellation_token: Optional token for cancelling operation + """ + self._contents.append(str(content)) + + async def query( + self, query: Union[str, MemoryContent], cancellation_token: Optional[CancellationToken] = None, **kwargs: Any + ) -> MemoryQueryResultSet: + """Query memory contents. + + Args: + query: Search query + cancellation_token: Optional token for cancelling operation + kwargs: Additional query parameters + + Returns: + Query results containing all memory contents + """ + results = [MemoryContent(content=content, mime_type="text/plain") for content in self._contents] + return MemoryQueryResultSet(results=results) + + async def clear(self, cancellation_token: Optional[CancellationToken] = None) -> None: + """Clear all memory contents. + + Args: + cancellation_token: Optional token for cancelling operation + """ + self._contents.clear() + + async def close(self) -> None: + """Close memory resources.""" + pass + + async def update_context(self, model_context: Any) -> UpdateContextResult: + """Update model context with memory contents. + + Args: + model_context: Context to update + + Returns: + Update result containing memory contents + """ + if self._contents: + results = [MemoryContent(content=content, mime_type="text/plain") for content in self._contents] + return UpdateContextResult(memories=MemoryQueryResultSet(results=results)) + return UpdateContextResult(memories=MemoryQueryResultSet(results=[])) + + def dump_component(self) -> ComponentModel: + """Dump memory state as component model. + + Returns: + Component model representing memory state + """ + return ComponentModel(provider="test", config={"type": "mock_memory"}) + + +class StructuredOutput(BaseModel): + """Test structured output model. + + Attributes: + content: Main content string + confidence: Confidence score between 0 and 1 + """ + + content: str + confidence: float + + @pytest.mark.asyncio -async def test_run_with_tools(monkeypatch: pytest.MonkeyPatch) -> None: - model_client = ReplayChatCompletionClient( +async def test_model_client_stream() -> None: + mock_client = ReplayChatCompletionClient( [ - CreateResult( - finish_reason="function_calls", - content=[FunctionCall(id="1", arguments=json.dumps({"input": "task"}), name="_pass_function")], - usage=RequestUsage(prompt_tokens=10, completion_tokens=5), - thought="Calling pass function", - cached=False, - ), - "pass", - "TERMINATE", - ], - model_info={ - "function_calling": True, - "vision": True, - "json_output": True, - "family": ModelFamily.GPT_4O, - "structured_output": True, - }, + "Response to message 3", + ] ) agent = AssistantAgent( - "tool_use_agent", - model_client=model_client, - tools=[ - _pass_function, - _fail_function, - FunctionTool(_echo_function, description="Echo"), - ], + "test_agent", + model_client=mock_client, + model_client_stream=True, ) - result = await agent.run(task="task") - - # Make sure the create call was made with the correct parameters. - assert len(model_client.create_calls) == 1 - llm_messages = model_client.create_calls[0]["messages"] - assert len(llm_messages) == 2 - assert isinstance(llm_messages[0], SystemMessage) - assert llm_messages[0].content == agent._system_messages[0].content # type: ignore - assert isinstance(llm_messages[1], UserMessage) - assert llm_messages[1].content == "task" - - assert len(result.messages) == 5 - assert isinstance(result.messages[0], TextMessage) - assert result.messages[0].models_usage is None - assert isinstance(result.messages[1], ThoughtEvent) - assert result.messages[1].content == "Calling pass function" - assert isinstance(result.messages[2], ToolCallRequestEvent) - assert result.messages[2].models_usage is not None - assert result.messages[2].models_usage.completion_tokens == 5 - assert result.messages[2].models_usage.prompt_tokens == 10 - assert isinstance(result.messages[3], ToolCallExecutionEvent) - assert result.messages[3].models_usage is None - assert isinstance(result.messages[4], ToolCallSummaryMessage) - assert result.messages[4].content == "pass" - assert result.messages[4].models_usage is None - - # Test streaming. - model_client.reset() - index = 0 + chunks: List[str] = [] async for message in agent.run_stream(task="task"): if isinstance(message, TaskResult): - assert message == result - else: - assert message == result.messages[index] - index += 1 - - # Test state saving and loading. - state = await agent.save_state() - agent2 = AssistantAgent( - "tool_use_agent", - model_client=model_client, - tools=[_pass_function, _fail_function, FunctionTool(_echo_function, description="Echo")], - ) - await agent2.load_state(state) - state2 = await agent2.save_state() - assert state == state2 + assert isinstance(message.messages[-1], TextMessage) + assert message.messages[-1].content == "Response to message 3" + elif isinstance(message, ModelClientStreamingChunkEvent): + chunks.append(message.content) + assert "".join(chunks) == "Response to message 3" @pytest.mark.asyncio -async def test_run_with_tools_and_reflection() -> None: - model_client = ReplayChatCompletionClient( +async def test_model_client_stream_with_tool_calls() -> None: + mock_client = ReplayChatCompletionClient( [ CreateResult( + content=[ + FunctionCall(id="1", name="_pass_function", arguments=r'{"input": "task"}'), + FunctionCall(id="3", name="_echo_function", arguments=r'{"input": "task"}'), + ], finish_reason="function_calls", - content=[FunctionCall(id="1", arguments=json.dumps({"input": "task"}), name="_pass_function")], - usage=RequestUsage(prompt_tokens=10, completion_tokens=5), - cached=False, - ), - CreateResult( - finish_reason="stop", - content="Hello", - usage=RequestUsage(prompt_tokens=10, completion_tokens=5), - cached=False, - ), - CreateResult( - finish_reason="stop", - content="TERMINATE", usage=RequestUsage(prompt_tokens=10, completion_tokens=5), cached=False, ), - ], - model_info={ - "function_calling": True, - "vision": True, - "json_output": True, - "family": ModelFamily.GPT_4O, - "structured_output": True, - }, + "Example response 2 to task", + ] ) + mock_client._model_info["function_calling"] = True # pyright: ignore agent = AssistantAgent( - "tool_use_agent", - model_client=model_client, - tools=[_pass_function, _fail_function, FunctionTool(_echo_function, description="Echo")], + "test_agent", + model_client=mock_client, + model_client_stream=True, reflect_on_tool_use=True, + tools=[_pass_function, _echo_function], ) - result = await agent.run(task="task") - - # Make sure the create call was made with the correct parameters. - assert len(model_client.create_calls) == 2 - llm_messages = model_client.create_calls[0]["messages"] - assert len(llm_messages) == 2 - assert isinstance(llm_messages[0], SystemMessage) - assert llm_messages[0].content == agent._system_messages[0].content # type: ignore - assert isinstance(llm_messages[1], UserMessage) - assert llm_messages[1].content == "task" - llm_messages = model_client.create_calls[1]["messages"] - assert len(llm_messages) == 4 - assert isinstance(llm_messages[0], SystemMessage) - assert llm_messages[0].content == agent._system_messages[0].content # type: ignore - assert isinstance(llm_messages[1], UserMessage) - assert llm_messages[1].content == "task" - assert isinstance(llm_messages[2], AssistantMessage) - assert isinstance(llm_messages[3], FunctionExecutionResultMessage) - - assert len(result.messages) == 4 - assert isinstance(result.messages[0], TextMessage) - assert result.messages[0].models_usage is None - assert isinstance(result.messages[1], ToolCallRequestEvent) - assert result.messages[1].models_usage is not None - assert result.messages[1].models_usage.completion_tokens == 5 - assert result.messages[1].models_usage.prompt_tokens == 10 - assert isinstance(result.messages[2], ToolCallExecutionEvent) - assert result.messages[2].models_usage is None - assert isinstance(result.messages[3], TextMessage) - assert result.messages[3].content == "Hello" - assert result.messages[3].models_usage is not None - assert result.messages[3].models_usage.completion_tokens == 5 - assert result.messages[3].models_usage.prompt_tokens == 10 - - # Test streaming. - model_client.reset() - index = 0 + chunks: List[str] = [] async for message in agent.run_stream(task="task"): if isinstance(message, TaskResult): - assert message == result - else: - assert message == result.messages[index] - index += 1 - - # Test state saving and loading. - state = await agent.save_state() - agent2 = AssistantAgent( - "tool_use_agent", - model_client=model_client, - tools=[ - _pass_function, - _fail_function, - FunctionTool(_echo_function, description="Echo"), - ], - ) - await agent2.load_state(state) - state2 = await agent2.save_state() - assert state == state2 + assert isinstance(message.messages[-1], TextMessage) + assert isinstance(message.messages[1], ToolCallRequestEvent) + assert message.messages[-1].content == "Example response 2 to task" + assert message.messages[1].content == [ + FunctionCall(id="1", name="_pass_function", arguments=r'{"input": "task"}'), + FunctionCall(id="3", name="_echo_function", arguments=r'{"input": "task"}'), + ] + assert isinstance(message.messages[2], ToolCallExecutionEvent) + assert message.messages[2].content == [ + FunctionExecutionResult(call_id="1", content="pass", is_error=False, name="_pass_function"), + FunctionExecutionResult(call_id="3", content="task", is_error=False, name="_echo_function"), + ] + elif isinstance(message, ModelClientStreamingChunkEvent): + chunks.append(message.content) + assert "".join(chunks) == "Example response 2 to task" @pytest.mark.asyncio -async def test_run_with_parallel_tools() -> None: +async def test_invalid_structured_output_format() -> None: + class AgentResponse(BaseModel): + response: str + status: str + model_client = ReplayChatCompletionClient( [ CreateResult( - finish_reason="function_calls", - content=[ - FunctionCall(id="1", arguments=json.dumps({"input": "task1"}), name="_pass_function"), - FunctionCall(id="2", arguments=json.dumps({"input": "task2"}), name="_pass_function"), - FunctionCall(id="3", arguments=json.dumps({"input": "task3"}), name="_echo_function"), - ], + finish_reason="stop", + content='{"response": "Hello"}', usage=RequestUsage(prompt_tokens=10, completion_tokens=5), - thought="Calling pass and echo functions", cached=False, ), - "pass", - "TERMINATE", - ], - model_info={ - "function_calling": True, - "vision": True, - "json_output": True, - "family": ModelFamily.GPT_4O, - "structured_output": True, - }, + ] ) + agent = AssistantAgent( - "tool_use_agent", - model_client=model_client, - tools=[ - _pass_function, - _fail_function, - FunctionTool(_echo_function, description="Echo"), - ], - ) - result = await agent.run(task="task") - - assert len(result.messages) == 5 - assert isinstance(result.messages[0], TextMessage) - assert result.messages[0].models_usage is None - assert isinstance(result.messages[1], ThoughtEvent) - assert result.messages[1].content == "Calling pass and echo functions" - assert isinstance(result.messages[2], ToolCallRequestEvent) - assert result.messages[2].content == [ - FunctionCall(id="1", arguments=r'{"input": "task1"}', name="_pass_function"), - FunctionCall(id="2", arguments=r'{"input": "task2"}', name="_pass_function"), - FunctionCall(id="3", arguments=r'{"input": "task3"}', name="_echo_function"), - ] - assert result.messages[2].models_usage is not None - assert result.messages[2].models_usage.completion_tokens == 5 - assert result.messages[2].models_usage.prompt_tokens == 10 - assert isinstance(result.messages[3], ToolCallExecutionEvent) - expected_content = [ - FunctionExecutionResult(call_id="1", content="pass", is_error=False, name="_pass_function"), - FunctionExecutionResult(call_id="2", content="pass", is_error=False, name="_pass_function"), - FunctionExecutionResult(call_id="3", content="task3", is_error=False, name="_echo_function"), - ] - for expected in expected_content: - assert expected in result.messages[3].content - assert result.messages[3].models_usage is None - assert isinstance(result.messages[4], ToolCallSummaryMessage) - assert result.messages[4].content == "pass\npass\ntask3" - assert result.messages[4].models_usage is None - - # Test streaming. - model_client.reset() - index = 0 - async for message in agent.run_stream(task="task"): - if isinstance(message, TaskResult): - assert message == result - else: - assert message == result.messages[index] - index += 1 - - # Test state saving and loading. - state = await agent.save_state() - agent2 = AssistantAgent( - "tool_use_agent", + name="assistant", model_client=model_client, - tools=[_pass_function, _fail_function, FunctionTool(_echo_function, description="Echo")], + output_content_type=AgentResponse, ) - await agent2.load_state(state) - state2 = await agent2.save_state() - assert state == state2 + + with pytest.raises(ValidationError): + await agent.run() @pytest.mark.asyncio -async def test_run_with_parallel_tools_with_empty_call_ids() -> None: +async def test_structured_message_factory_serialization() -> None: + class AgentResponse(BaseModel): + result: str + status: str + model_client = ReplayChatCompletionClient( [ CreateResult( - finish_reason="function_calls", - content=[ - FunctionCall(id="", arguments=json.dumps({"input": "task1"}), name="_pass_function"), - FunctionCall(id="", arguments=json.dumps({"input": "task2"}), name="_pass_function"), - FunctionCall(id="", arguments=json.dumps({"input": "task3"}), name="_echo_function"), - ], + finish_reason="stop", + content=AgentResponse(result="All good", status="ok").model_dump_json(), usage=RequestUsage(prompt_tokens=10, completion_tokens=5), cached=False, - ), - "pass", - "TERMINATE", - ], - model_info={ - "function_calling": True, - "vision": True, - "json_output": True, - "family": ModelFamily.GPT_4O, - "structured_output": True, - }, + ) + ] ) + agent = AssistantAgent( - "tool_use_agent", + name="structured_agent", model_client=model_client, - tools=[ - _pass_function, - _fail_function, - FunctionTool(_echo_function, description="Echo"), - ], + output_content_type=AgentResponse, + output_content_type_format="{result} - {status}", ) - result = await agent.run(task="task") - - assert len(result.messages) == 4 - assert isinstance(result.messages[0], TextMessage) - assert result.messages[0].models_usage is None - assert isinstance(result.messages[1], ToolCallRequestEvent) - assert result.messages[1].content == [ - FunctionCall(id="", arguments=r'{"input": "task1"}', name="_pass_function"), - FunctionCall(id="", arguments=r'{"input": "task2"}', name="_pass_function"), - FunctionCall(id="", arguments=r'{"input": "task3"}', name="_echo_function"), - ] - assert result.messages[1].models_usage is not None - assert result.messages[1].models_usage.completion_tokens == 5 - assert result.messages[1].models_usage.prompt_tokens == 10 - assert isinstance(result.messages[2], ToolCallExecutionEvent) - expected_content = [ - FunctionExecutionResult(call_id="", content="pass", is_error=False, name="_pass_function"), - FunctionExecutionResult(call_id="", content="pass", is_error=False, name="_pass_function"), - FunctionExecutionResult(call_id="", content="task3", is_error=False, name="_echo_function"), - ] - for expected in expected_content: - assert expected in result.messages[2].content - assert result.messages[2].models_usage is None - assert isinstance(result.messages[3], ToolCallSummaryMessage) - assert result.messages[3].content == "pass\npass\ntask3" - assert result.messages[3].models_usage is None - - # Test streaming. - model_client.reset() - index = 0 - async for message in agent.run_stream(task="task"): - if isinstance(message, TaskResult): - assert message == result - else: - assert message == result.messages[index] - index += 1 - - # Test state saving and loading. - state = await agent.save_state() - agent2 = AssistantAgent( - "tool_use_agent", - model_client=model_client, - tools=[_pass_function, _fail_function, FunctionTool(_echo_function, description="Echo")], - ) - await agent2.load_state(state) - state2 = await agent2.save_state() - assert state == state2 + + dumped = agent.dump_component() + restored_agent = AssistantAgent.load_component(dumped) + result = await restored_agent.run() + + assert isinstance(result.messages[0], StructuredMessage) + assert result.messages[0].content.result == "All good" # type: ignore + assert result.messages[0].content.status == "ok" # type: ignore @pytest.mark.asyncio -async def test_handoffs() -> None: - handoff = Handoff(target="agent2") +async def test_structured_message_format_string() -> None: + class AgentResponse(BaseModel): + field1: str + field2: str + + expected = AgentResponse(field1="foo", field2="bar") + model_client = ReplayChatCompletionClient( [ CreateResult( - finish_reason="function_calls", - content=[ - FunctionCall(id="1", arguments=json.dumps({}), name=handoff.name), - ], - usage=RequestUsage(prompt_tokens=42, completion_tokens=43), + finish_reason="stop", + content=expected.model_dump_json(), + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), cached=False, ) - ], - model_info={ - "function_calling": True, - "vision": True, - "json_output": True, - "family": ModelFamily.GPT_4O, - "structured_output": True, - }, + ] ) - tool_use_agent = AssistantAgent( - "tool_use_agent", + + agent = AssistantAgent( + name="formatted_agent", model_client=model_client, - tools=[ - _pass_function, - _fail_function, - FunctionTool(_echo_function, description="Echo"), - ], - handoffs=[handoff], + output_content_type=AgentResponse, + output_content_type_format="{field1} - {field2}", ) - assert HandoffMessage in tool_use_agent.produced_message_types - result = await tool_use_agent.run(task="task") - assert len(result.messages) == 4 - assert isinstance(result.messages[0], TextMessage) - assert result.messages[0].models_usage is None - assert isinstance(result.messages[1], ToolCallRequestEvent) - assert result.messages[1].models_usage is not None - assert result.messages[1].models_usage.completion_tokens == 43 - assert result.messages[1].models_usage.prompt_tokens == 42 - assert isinstance(result.messages[2], ToolCallExecutionEvent) - assert result.messages[2].models_usage is None - assert isinstance(result.messages[3], HandoffMessage) - assert result.messages[3].content == handoff.message - assert result.messages[3].target == handoff.target - assert result.messages[3].models_usage is None - - # Test streaming. - model_client.reset() - index = 0 - async for message in tool_use_agent.run_stream(task="task"): - if isinstance(message, TaskResult): - assert message == result - else: - assert message == result.messages[index] - index += 1 + result = await agent.run() -@pytest.mark.asyncio -async def test_custom_handoffs() -> None: - name = "transfer_to_agent2" - description = "Handoff to agent2." - next_action = "next_action" + assert len(result.messages) == 1 + message = result.messages[0] - class TextCommandHandOff(Handoff): - @property - def handoff_tool(self) -> BaseTool[BaseModel, BaseModel]: - """Create a handoff tool from this handoff configuration.""" + # Check that it's a StructuredMessage with the correct content model + assert isinstance(message, StructuredMessage) + assert isinstance(message.content, AgentResponse) # type: ignore[reportUnknownMemberType] + assert message.content == expected - def _next_action(action: str) -> str: - """Returns the action you want the user to perform""" - return action + # Check that the format_string was applied correctly + assert message.to_model_text() == "foo - bar" - return FunctionTool(_next_action, name=self.name, description=self.description, strict=True) - handoff = TextCommandHandOff(name=name, description=description, target="agent2") - model_client = ReplayChatCompletionClient( - [ - CreateResult( - finish_reason="function_calls", - content=[ - FunctionCall(id="1", arguments=json.dumps({"action": next_action}), name=handoff.name), - ], - usage=RequestUsage(prompt_tokens=42, completion_tokens=43), - cached=False, - ) - ], - model_info={ - "function_calling": True, - "vision": True, - "json_output": True, - "family": ModelFamily.GPT_4O, - "structured_output": True, - }, - ) - tool_use_agent = AssistantAgent( - "tool_use_agent", - model_client=model_client, - tools=[ - _pass_function, - _fail_function, - FunctionTool(_echo_function, description="Echo"), - ], - handoffs=[handoff], +@pytest.mark.asyncio +async def test_tools_serialize_and_deserialize() -> None: + def test() -> str: + return "hello world" + + client = OpenAIChatCompletionClient( + model="gpt-4o", + api_key="API_KEY", ) - assert HandoffMessage in tool_use_agent.produced_message_types - result = await tool_use_agent.run(task="task") - assert len(result.messages) == 4 - assert isinstance(result.messages[0], TextMessage) - assert result.messages[0].models_usage is None - assert isinstance(result.messages[1], ToolCallRequestEvent) - assert result.messages[1].models_usage is not None - assert result.messages[1].models_usage.completion_tokens == 43 - assert result.messages[1].models_usage.prompt_tokens == 42 - assert isinstance(result.messages[2], ToolCallExecutionEvent) - assert result.messages[2].models_usage is None - assert isinstance(result.messages[3], HandoffMessage) - assert result.messages[3].content == next_action - assert result.messages[3].target == handoff.target - - assert result.messages[3].models_usage is None - - # Test streaming. - model_client.reset() - index = 0 - async for message in tool_use_agent.run_stream(task="task"): - if isinstance(message, TaskResult): - assert message == result - else: - assert message == result.messages[index] - index += 1 + agent = AssistantAgent( + name="test", + model_client=client, + tools=[test], + ) -@pytest.mark.asyncio -async def test_custom_object_handoffs() -> None: - """test handoff tool return a object""" - name = "transfer_to_agent2" - description = "Handoff to agent2." - next_action = {"action": "next_action"} # using a map, not a str + serialize = agent.dump_component() + deserialize = AssistantAgent.load_component(serialize) - class DictCommandHandOff(Handoff): - @property - def handoff_tool(self) -> BaseTool[BaseModel, BaseModel]: - """Create a handoff tool from this handoff configuration.""" + assert deserialize.name == agent.name + for original, restored in zip(agent._workbench, deserialize._workbench, strict=True): # type: ignore + assert await original.list_tools() == await restored.list_tools() # type: ignore + assert agent.component_version == deserialize.component_version - def _next_action(action: str) -> Dict[str, str]: - """Returns the action you want the user to perform""" - return {"action": action} - return FunctionTool(_next_action, name=self.name, description=self.description, strict=True) +@pytest.mark.asyncio +async def test_workbench_serialize_and_deserialize() -> None: + workbench = McpWorkbench(server_params=SseServerParams(url="http://test-url")) - handoff = DictCommandHandOff(name=name, description=description, target="agent2") - model_client = ReplayChatCompletionClient( - [ - CreateResult( - finish_reason="function_calls", - content=[ - FunctionCall(id="1", arguments=json.dumps({"action": "next_action"}), name=handoff.name), - ], - usage=RequestUsage(prompt_tokens=42, completion_tokens=43), - cached=False, - ) - ], - model_info={ - "function_calling": True, - "vision": True, - "json_output": True, - "family": ModelFamily.GPT_4O, - "structured_output": True, - }, + client = OpenAIChatCompletionClient( + model="gpt-4o", + api_key="API_KEY", ) - tool_use_agent = AssistantAgent( - "tool_use_agent", - model_client=model_client, - tools=[ - _pass_function, - _fail_function, - FunctionTool(_echo_function, description="Echo"), - ], - handoffs=[handoff], + + agent = AssistantAgent( + name="test", + model_client=client, + workbench=workbench, ) - assert HandoffMessage in tool_use_agent.produced_message_types - result = await tool_use_agent.run(task="task") - assert len(result.messages) == 4 - assert isinstance(result.messages[0], TextMessage) - assert result.messages[0].models_usage is None - assert isinstance(result.messages[1], ToolCallRequestEvent) - assert result.messages[1].models_usage is not None - assert result.messages[1].models_usage.completion_tokens == 43 - assert result.messages[1].models_usage.prompt_tokens == 42 - assert isinstance(result.messages[2], ToolCallExecutionEvent) - assert result.messages[2].models_usage is None - assert isinstance(result.messages[3], HandoffMessage) - # the content will return as a string, because the function call will convert to string - assert result.messages[3].content == str(next_action) - assert result.messages[3].target == handoff.target - - assert result.messages[3].models_usage is None - - # Test streaming. - model_client.reset() - index = 0 - async for message in tool_use_agent.run_stream(task="task"): - if isinstance(message, TaskResult): - assert message == result - else: - assert message == result.messages[index] - index += 1 + + serialize = agent.dump_component() + deserialize = AssistantAgent.load_component(serialize) + + assert deserialize.name == agent.name + for original, restored in zip(agent._workbench, deserialize._workbench, strict=True): # type: ignore + assert isinstance(original, McpWorkbench) + assert isinstance(restored, McpWorkbench) + assert original._to_config() == restored._to_config() # type: ignore @pytest.mark.asyncio -async def test_multi_modal_task(monkeypatch: pytest.MonkeyPatch) -> None: - model_client = ReplayChatCompletionClient(["Hello"]) +async def test_multiple_workbenches_serialize_and_deserialize() -> None: + workbenches: List[McpWorkbench] = [ + McpWorkbench(server_params=SseServerParams(url="http://test-url-1")), + McpWorkbench(server_params=SseServerParams(url="http://test-url-2")), + ] + + client = OpenAIChatCompletionClient( + model="gpt-4o", + api_key="API_KEY", + ) + agent = AssistantAgent( - name="assistant", - model_client=model_client, + name="test_multi", + model_client=client, + workbench=workbenches, ) - # Generate a random base64 image. - img_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC" - result = await agent.run(task=MultiModalMessage(source="user", content=["Test", Image.from_base64(img_base64)])) - assert len(result.messages) == 2 + + serialize = agent.dump_component() + deserialized_agent: AssistantAgent = AssistantAgent.load_component(serialize) + + assert deserialized_agent.name == agent.name + assert isinstance(deserialized_agent._workbench, list) # type: ignore + assert len(deserialized_agent._workbench) == len(workbenches) # type: ignore + + for original, restored in zip(agent._workbench, deserialized_agent._workbench, strict=True): # type: ignore + assert isinstance(original, McpWorkbench) + assert isinstance(restored, McpWorkbench) + assert original._to_config() == restored._to_config() # type: ignore @pytest.mark.asyncio -async def test_invalid_model_capabilities() -> None: - model = "random-model" - model_client = OpenAIChatCompletionClient( - model=model, - api_key="", - model_info={ - "vision": False, - "function_calling": False, - "json_output": False, - "family": ModelFamily.UNKNOWN, - "structured_output": False, - }, - ) +async def test_tools_deserialize_aware() -> None: + dump = """ + { + "provider": "autogen_agentchat.agents.AssistantAgent", + "component_type": "agent", + "version": 1, + "component_version": 2, + "description": "An agent that provides assistance with tool use.", + "label": "AssistantAgent", + "config": { + "name": "TestAgent", + "model_client":{ + "provider": "autogen_ext.models.replay.ReplayChatCompletionClient", + "component_type": "replay_chat_completion_client", + "version": 1, + "component_version": 1, + "description": "A mock chat completion client that replays predefined responses using an index-based approach.", + "label": "ReplayChatCompletionClient", + "config": { + "chat_completions": [ + { + "finish_reason": "function_calls", + "content": [ + { + "id": "hello", + "arguments": "{}", + "name": "hello" + } + ], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0 + }, + "cached": false + } + ], + "model_info": { + "vision": false, + "function_calling": true, + "json_output": false, + "family": "unknown", + "structured_output": false + } + } + }, + "tools": [ + { + "provider": "autogen_core.tools.FunctionTool", + "component_type": "tool", + "version": 1, + "component_version": 1, + "description": "Create custom tools by wrapping standard Python functions.", + "label": "FunctionTool", + "config": { + "source_code": "def hello():\\n return 'Hello, World!'\\n", + "name": "hello", + "description": "", + "global_imports": [], + "has_cancellation_support": false + } + } + ], + "model_context": { + "provider": "autogen_core.model_context.UnboundedChatCompletionContext", + "component_type": "chat_completion_context", + "version": 1, + "component_version": 1, + "description": "An unbounded chat completion context that keeps a view of the all the messages.", + "label": "UnboundedChatCompletionContext", + "config": {} + }, + "description": "An agent that provides assistance with ability to use tools.", + "system_message": "You are a helpful assistant.", + "model_client_stream": false, + "reflect_on_tool_use": false, + "tool_call_summary_format": "{result}", + "metadata": {} + } + } + + """ + + # Test that agent can be deserialized from configuration + config = json.loads(dump) + agent = AssistantAgent.load_component(config) + + # Verify the agent was loaded correctly + assert agent.name == "TestAgent" + assert agent.description == "An agent that provides assistance with ability to use tools." + + +class TestAssistantAgentToolCallLoop: + """Test suite for tool call loop functionality. + + Tests the behavior of AssistantAgent's tool call loop feature, which allows + multiple sequential tool calls before producing a final response. + """ + + @pytest.mark.asyncio + async def test_tool_call_loop_enabled(self) -> None: + """Test that tool call loop works when enabled. + + Verifies that: + 1. Multiple tool calls are executed in sequence + 2. Loop continues until non-tool response + 3. Final response is correct type + """ + # Create mock client with multiple tool calls followed by text response + model_client = ReplayChatCompletionClient( + [ + # First tool call + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "first"}), name="mock_tool_function")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + # Second tool call (loop continues) + CreateResult( + finish_reason="function_calls", + content=[ + FunctionCall(id="2", arguments=json.dumps({"param": "second"}), name="mock_tool_function") + ], + usage=RequestUsage(prompt_tokens=12, completion_tokens=5), + cached=False, + ), + # Final text response (loop ends) + CreateResult( + finish_reason="stop", + content="Task completed successfully!", + usage=RequestUsage(prompt_tokens=15, completion_tokens=10), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) - with pytest.raises(ValueError): agent = AssistantAgent( - name="assistant", + name="test_agent", model_client=model_client, - tools=[ - _pass_function, - _fail_function, - FunctionTool(_echo_function, description="Echo"), - ], + tools=[mock_tool_function], + max_tool_iterations=3, ) - await agent.run(task=TextMessage(source="user", content="Test")) - with pytest.raises(ValueError): - agent = AssistantAgent(name="assistant", model_client=model_client, handoffs=["agent2"]) - await agent.run(task=TextMessage(source="user", content="Test")) + result = await agent.run(task="Execute multiple tool calls") + + # Verify multiple model calls were made + assert len(model_client.create_calls) == 3, f"Expected 3 calls, got {len(model_client.create_calls)}" + + # Verify final response is text + final_message = result.messages[-1] + assert isinstance(final_message, TextMessage) + assert final_message.content == "Task completed successfully!" + + @pytest.mark.asyncio + async def test_tool_call_loop_disabled_default(self) -> None: + """Test that tool call loop is disabled by default. + + Verifies that: + 1. Only one tool call is made when loop is disabled + 2. Agent returns after first tool call + """ + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="mock_tool_function")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + max_tool_iterations=1, + ) -@pytest.mark.asyncio -async def test_remove_images() -> None: - model = "random-model" - model_client_1 = OpenAIChatCompletionClient( - model=model, - api_key="", - model_info={ - "vision": False, - "function_calling": False, - "json_output": False, - "family": ModelFamily.UNKNOWN, - "structured_output": False, - }, - ) - model_client_2 = OpenAIChatCompletionClient( - model=model, - api_key="", - model_info={ - "vision": True, - "function_calling": False, - "json_output": False, - "family": ModelFamily.UNKNOWN, - "structured_output": False, - }, - ) + result = await agent.run(task="Execute single tool call") + + # Should only make one model call + assert len(model_client.create_calls) == 1, f"Expected 1 call, got {len(model_client.create_calls)}" + assert result is not None + + @pytest.mark.asyncio + async def test_tool_call_loop_max_iterations(self) -> None: + """Test that tool call loop respects max_iterations limit.""" + # Create responses that would continue forever without max_iterations + responses: List[CreateResult] = [] + for i in range(15): # More than default max_iterations (10) + responses.append( + CreateResult( + finish_reason="function_calls", + content=[ + FunctionCall(id=str(i), arguments=json.dumps({"param": f"call_{i}"}), name="mock_tool_function") + ], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ) - img_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC" - messages: List[LLMMessage] = [ - SystemMessage(content="System.1"), - UserMessage(content=["User.1", Image.from_base64(img_base64)], source="user.1"), - AssistantMessage(content="Assistant.1", source="assistant.1"), - UserMessage(content="User.2", source="assistant.2"), - ] + model_client = ReplayChatCompletionClient( + responses, + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) - agent_1 = AssistantAgent(name="assistant_1", model_client=model_client_1) - result = agent_1._get_compatible_context(model_client_1, messages) # type: ignore - assert len(result) == 4 - assert isinstance(result[1].content, str) + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + max_tool_iterations=5, # Set max iterations to 5 + ) - agent_2 = AssistantAgent(name="assistant_2", model_client=model_client_2) - result = agent_2._get_compatible_context(model_client_2, messages) # type: ignore - assert len(result) == 4 - assert isinstance(result[1].content, list) + result = await agent.run(task="Test max iterations") + + # Should stop at max_iterations + assert len(model_client.create_calls) == 5, f"Expected 5 calls, got {len(model_client.create_calls)}" + # Verify result is not None + assert result is not None + + @pytest.mark.asyncio + async def test_tool_call_loop_with_handoff(self) -> None: + """Test that tool call loop stops on handoff.""" + model_client = ReplayChatCompletionClient( + [ + # Tool call followed by handoff + CreateResult( + finish_reason="function_calls", + content=[ + FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="mock_tool_function"), + FunctionCall( + id="2", arguments=json.dumps({"target": "other_agent"}), name="transfer_to_other_agent" + ), + ], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + handoffs=["other_agent"], + max_tool_iterations=1, + ) -@pytest.mark.asyncio -async def test_list_chat_messages(monkeypatch: pytest.MonkeyPatch) -> None: - model_client = ReplayChatCompletionClient( - [ - CreateResult( - finish_reason="stop", - content="Response to message 1", - usage=RequestUsage(prompt_tokens=10, completion_tokens=5), - cached=False, + result = await agent.run(task="Test handoff in loop") + + # Should stop at handoff + assert len(model_client.create_calls) == 1, f"Expected 1 call, got {len(model_client.create_calls)}" + + # Should return HandoffMessage + assert isinstance(result.messages[-1], HandoffMessage) + + @pytest.mark.asyncio + async def test_tool_call_config_validation(self) -> None: + """Test that ToolCallConfig validation works correctly.""" + # Test that max_iterations must be >= 1 + with pytest.raises( + ValueError, match="Maximum number of tool iterations must be greater than or equal to 1, got 0" + ): + AssistantAgent( + name="test_agent", + model_client=MagicMock(), + max_tool_iterations=0, # Should raise error ) - ] - ) - agent = AssistantAgent( - "test_agent", - model_client=model_client, - ) - - # Create a list of chat messages - messages: List[ChatMessage] = [ - TextMessage(content="Message 1", source="user"), - TextMessage(content="Message 2", source="user"), - ] - # Test run method with list of messages - result = await agent.run(task=messages) - assert len(result.messages) == 3 # 2 input messages + 1 response message - assert isinstance(result.messages[0], TextMessage) - assert result.messages[0].content == "Message 1" - assert result.messages[0].source == "user" - assert isinstance(result.messages[1], TextMessage) - assert result.messages[1].content == "Message 2" - assert result.messages[1].source == "user" - assert isinstance(result.messages[2], TextMessage) - assert result.messages[2].content == "Response to message 1" - assert result.messages[2].source == "test_agent" - assert result.messages[2].models_usage is not None - assert result.messages[2].models_usage.completion_tokens == 5 - assert result.messages[2].models_usage.prompt_tokens == 10 - - # Test run_stream method with list of messages - model_client.reset() # Reset the mock client - index = 0 - async for message in agent.run_stream(task=messages): - if isinstance(message, TaskResult): - assert message == result - else: - assert message == result.messages[index] - index += 1 +class TestAssistantAgentInitialization: + """Test suite for AssistantAgent initialization. + + Tests various initialization scenarios and configurations of the AssistantAgent class. + """ + + @pytest.mark.asyncio + async def test_basic_initialization(self) -> None: + """Test basic agent initialization with minimal parameters. + + Verifies that: + 1. Agent initializes with required parameters + 2. Default values are set correctly + 3. Basic functionality works + """ + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="stop", + content="Hello!", + usage=RequestUsage(prompt_tokens=5, completion_tokens=2), + cached=False, + ) + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) -@pytest.mark.asyncio -async def test_model_context(monkeypatch: pytest.MonkeyPatch) -> None: - model_client = ReplayChatCompletionClient(["Response to message 3"]) - model_context = BufferedChatCompletionContext(buffer_size=2) - agent = AssistantAgent( - "test_agent", - model_client=model_client, - model_context=model_context, - ) + agent = AssistantAgent(name="test_agent", model_client=model_client) + result = await agent.run(task="Say hello") + + assert isinstance(result.messages[-1], TextMessage) + assert result.messages[-1].content == "Hello!" + + @pytest.mark.asyncio + async def test_initialization_with_tools(self) -> None: + """Test agent initialization with tools. + + Verifies that: + 1. Agent accepts tool configurations + 2. Tools are properly registered + 3. Tool calls work correctly + """ + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="mock_tool_function")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) - messages = [ - TextMessage(content="Message 1", source="user"), - TextMessage(content="Message 2", source="user"), - TextMessage(content="Message 3", source="user"), - ] - await agent.run(task=messages) + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + ) - # Check that the model_context property returns the correct internal context - assert agent.model_context == model_context - # Check if the mock client is called with only the last two messages. - assert len(model_client.create_calls) == 1 - # 2 message from the context + 1 system message - assert len(model_client.create_calls[0]["messages"]) == 3 + result = await agent.run(task="Use the tool") + assert isinstance(result.messages[-1], ToolCallSummaryMessage) + assert "Tool executed with: test" in result.messages[-1].content + + @pytest.mark.asyncio + async def test_initialization_with_memory(self) -> None: + """Test agent initialization with memory. + + Verifies that: + 1. Memory is properly integrated + 2. Memory contents affect responses + 3. Memory updates work correctly + """ + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="stop", + content="Using memory content", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + memory = MockMemory(contents=["Test memory content"]) + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + memory=[memory], + ) -@pytest.mark.asyncio -async def test_run_with_memory(monkeypatch: pytest.MonkeyPatch) -> None: - model_client = ReplayChatCompletionClient(["Hello"]) - b64_image_str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC" - - # Test basic memory properties and empty context - memory = ListMemory(name="test_memory") - assert memory.name == "test_memory" - - empty_context = BufferedChatCompletionContext(buffer_size=2) - empty_results = await memory.update_context(empty_context) - assert len(empty_results.memories.results) == 0 - - # Test various content types - memory = ListMemory() - await memory.add(MemoryContent(content="text content", mime_type=MemoryMimeType.TEXT)) - await memory.add(MemoryContent(content={"key": "value"}, mime_type=MemoryMimeType.JSON)) - await memory.add(MemoryContent(content=Image.from_base64(b64_image_str), mime_type=MemoryMimeType.IMAGE)) - - # Test query functionality - query_result = await memory.query(MemoryContent(content="", mime_type=MemoryMimeType.TEXT)) - assert isinstance(query_result, MemoryQueryResult) - # Should have all three memories we added - assert len(query_result.results) == 3 - - # Test clear and cleanup - await memory.clear() - empty_query = await memory.query(MemoryContent(content="", mime_type=MemoryMimeType.TEXT)) - assert len(empty_query.results) == 0 - await memory.close() # Should not raise - - # Test invalid memory type - with pytest.raises(TypeError): - AssistantAgent( - "test_agent", - model_client=model_client, - memory="invalid", # type: ignore - ) - - # Test with agent - memory2 = ListMemory() - await memory2.add(MemoryContent(content="test instruction", mime_type=MemoryMimeType.TEXT)) - - agent = AssistantAgent("test_agent", model_client=model_client, memory=[memory2]) - - # Test dump and load component with memory - agent_config: ComponentModel = agent.dump_component() - assert agent_config.provider == "autogen_agentchat.agents.AssistantAgent" - agent2 = AssistantAgent.load_component(agent_config) - - result = await agent2.run(task="test task") - assert len(result.messages) > 0 - memory_event = next((msg for msg in result.messages if isinstance(msg, MemoryQueryEvent)), None) - assert memory_event is not None - assert len(memory_event.content) > 0 - assert isinstance(memory_event.content[0], MemoryContent) - - # Test memory protocol - class BadMemory: - pass + result = await agent.run(task="Use memory") + assert isinstance(result.messages[-1], TextMessage) + assert result.messages[-1].content == "Using memory content" - assert not isinstance(BadMemory(), Memory) - assert isinstance(ListMemory(), Memory) + @pytest.mark.asyncio + async def test_initialization_with_handoffs(self) -> None: + """Test agent initialization with handoffs.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": True, "vision": False, "family": ModelFamily.GPT_4O} + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + handoffs=["agent1", Handoff(target="agent2")], + ) -@pytest.mark.asyncio -async def test_assistant_agent_declarative() -> None: - model_client = ReplayChatCompletionClient( - ["Response to message 3"], - model_info={ - "function_calling": True, - "vision": True, - "json_output": True, - "family": ModelFamily.GPT_4O, - "structured_output": True, - }, - ) - model_context = BufferedChatCompletionContext(buffer_size=2) - agent = AssistantAgent( - "test_agent", - model_client=model_client, - model_context=model_context, - memory=[ListMemory(name="test_memory")], - ) + assert len(agent._handoffs) == 2 # type: ignore[reportPrivateUsage] + assert "transfer_to_agent1" in agent._handoffs # type: ignore[reportPrivateUsage] + assert "transfer_to_agent2" in agent._handoffs # type: ignore[reportPrivateUsage] - agent_config: ComponentModel = agent.dump_component() - assert agent_config.provider == "autogen_agentchat.agents.AssistantAgent" + @pytest.mark.asyncio + async def test_initialization_with_custom_model_context(self) -> None: + """Test agent initialization with custom model context.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} - agent2 = AssistantAgent.load_component(agent_config) - assert agent2.name == agent.name + model_context = BufferedChatCompletionContext(buffer_size=5) + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + model_context=model_context, + ) - agent3 = AssistantAgent( - "test_agent", - model_client=model_client, - model_context=model_context, - tools=[ - _pass_function, - _fail_function, - FunctionTool(_echo_function, description="Echo"), - ], - ) - agent3_config = agent3.dump_component() - assert agent3_config.provider == "autogen_agentchat.agents.AssistantAgent" + assert agent._model_context == model_context # type: ignore[reportPrivateUsage] + @pytest.mark.asyncio + async def test_initialization_with_structured_output(self) -> None: + """Test agent initialization with structured output.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} -@pytest.mark.asyncio -async def test_model_client_stream() -> None: - mock_client = ReplayChatCompletionClient( - [ - "Response to message 3", + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + output_content_type=StructuredOutput, + ) + + assert agent._output_content_type == StructuredOutput # type: ignore[reportPrivateUsage] + assert agent._reflect_on_tool_use is True # type: ignore[reportPrivateUsage] # Should be True by default with structured output + + @pytest.mark.asyncio + async def test_initialization_with_metadata(self) -> None: + """Test agent initialization with metadata.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + metadata = {"key1": "value1", "key2": "value2"} + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + metadata=metadata, + ) + + assert agent._metadata == metadata # type: ignore[reportPrivateUsage] + + @pytest.mark.asyncio + async def test_output_task_messages_false(self) -> None: + """Test agent with output_task_messages=False. + + Verifies that: + 1. Task messages are excluded from result when output_task_messages=False + 2. Only agent response messages are included in output + 3. Both run and run_stream respect the parameter + """ + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="stop", + content="Agent response without task message", + usage=RequestUsage(prompt_tokens=10, completion_tokens=8), + cached=False, + ), + CreateResult( + finish_reason="stop", + content="Second agent response", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": False, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent(name="test_agent", model_client=model_client) + + # Test run() with output_task_messages=False + result = await agent.run(task="Test task message", output_task_messages=False) + + # Should only contain the agent's response, not the task message + assert len(result.messages) == 1 + assert isinstance(result.messages[0], TextMessage) + assert result.messages[0].content == "Agent response without task message" + assert result.messages[0].source == "test_agent" # Test run_stream() with output_task_messages=False + # Create a new model client for streaming test to avoid response conflicts + stream_model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="stop", + content="Stream agent response", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": False, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + stream_agent = AssistantAgent(name="test_agent", model_client=stream_model_client) + streamed_messages: List[BaseAgentEvent | BaseChatMessage] = [] + final_result: TaskResult | None = None + + async for message in stream_agent.run_stream(task="Test task message", output_task_messages=False): + if isinstance(message, TaskResult): + final_result = message + else: + streamed_messages.append(message) + + # Verify streaming behavior + assert final_result is not None + assert len(final_result.messages) == 1 + assert isinstance(final_result.messages[0], TextMessage) + assert final_result.messages[0].content == "Stream agent response" + + # Verify that no task message was streamed + task_messages = [msg for msg in streamed_messages if isinstance(msg, TextMessage) and msg.source == "user"] + assert len(task_messages) == 0 # Test with multiple task messages + multi_model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="stop", + content="Multi task response", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": False, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + multi_agent = AssistantAgent(name="test_agent", model_client=multi_model_client) + task_messages_list = [ + TextMessage(content="First task", source="user"), + TextMessage(content="Second task", source="user"), ] - ) - agent = AssistantAgent( - "test_agent", - model_client=mock_client, - model_client_stream=True, - ) - chunks: List[str] = [] - async for message in agent.run_stream(task="task"): - if isinstance(message, TaskResult): - assert message.messages[-1].content == "Response to message 3" - elif isinstance(message, ModelClientStreamingChunkEvent): - chunks.append(message.content) - assert "".join(chunks) == "Response to message 3" + result_multi = await multi_agent.run(task=task_messages_list, output_task_messages=False) + + # Should only contain the agent's response, not the multiple task messages + assert len(result_multi.messages) == 1 + assert isinstance(result_multi.messages[0], TextMessage) + assert result_multi.messages[0].source == "test_agent" + assert result_multi.messages[0].content == "Multi task response" + + +class TestAssistantAgentValidation: + """Test suite for AssistantAgent validation. + + Tests various validation scenarios to ensure proper error handling and input validation. + """ + + @pytest.mark.asyncio + async def test_tool_names_must_be_unique(self) -> None: + """Test validation of unique tool names. + + Verifies that: + 1. Duplicate tool names are detected + 2. Appropriate error is raised + """ + + def duplicate_tool(param: str) -> str: + """Test tool with duplicate name. + + Args: + param: Input parameter + + Returns: + Formatted string with parameter + """ + return f"Duplicate tool: {param}" + + model_client = ReplayChatCompletionClient( + [], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + with pytest.raises(ValueError, match="Tool names must be unique"): + AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function, duplicate_tool, mock_tool_function], + ) + + @pytest.mark.asyncio + async def test_handoff_names_must_be_unique(self) -> None: + """Test validation of unique handoff names. + + Verifies that: + 1. Duplicate handoff names are detected + 2. Appropriate error is raised + """ + model_client = ReplayChatCompletionClient( + [], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + with pytest.raises(ValueError, match="Handoff names must be unique"): + AssistantAgent( + name="test_agent", + model_client=model_client, + handoffs=["agent1", "agent2", "agent1"], + ) + + @pytest.mark.asyncio + async def test_handoff_names_must_be_unique_from_tool_names(self) -> None: + """Test validation of handoff names against tool names. + + Verifies that: + 1. Handoff names cannot conflict with tool names + 2. Appropriate error is raised + """ + + def test_tool() -> str: + """Test tool with name that conflicts with handoff. + + Returns: + Static test string + """ + return "test" + + model_client = ReplayChatCompletionClient( + [], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + with pytest.raises(ValueError, match="Handoff names must be unique from tool names"): + AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[test_tool], + handoffs=["test_tool"], + ) + + @pytest.mark.asyncio + async def test_function_calling_required_for_tools(self) -> None: + """Test that function calling is required for tools.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + with pytest.raises(ValueError, match="The model does not support function calling"): + AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + ) + + @pytest.mark.asyncio + async def test_function_calling_required_for_handoffs(self) -> None: + """Test that function calling is required for handoffs.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + with pytest.raises( + ValueError, match="The model does not support function calling, which is needed for handoffs" + ): + AssistantAgent( + name="test_agent", + model_client=model_client, + handoffs=["agent1"], + ) + + @pytest.mark.asyncio + async def test_memory_type_validation(self) -> None: + """Test memory type validation.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + with pytest.raises(TypeError, match="Expected Memory, List\\[Memory\\], or None"): + AssistantAgent( + name="test_agent", + model_client=model_client, + memory="invalid_memory", # type: ignore + ) + + @pytest.mark.asyncio + async def test_tools_and_workbench_mutually_exclusive(self) -> None: + """Test that tools and workbench are mutually exclusive.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": True, "vision": False, "family": ModelFamily.GPT_4O} + + workbench = MagicMock() + + with pytest.raises(ValueError, match="Tools cannot be used with a workbench"): + AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + workbench=workbench, + ) + + @pytest.mark.asyncio + async def test_unsupported_tool_type(self) -> None: + """Test error handling for unsupported tool types.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": True, "vision": False, "family": ModelFamily.GPT_4O} + + with pytest.raises(ValueError, match="Unsupported tool type"): + AssistantAgent( + name="test_agent", + model_client=model_client, + tools=["invalid_tool"], # type: ignore + ) + + @pytest.mark.asyncio + async def test_unsupported_handoff_type(self) -> None: + """Test error handling for unsupported handoff types.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": True, "vision": False, "family": ModelFamily.GPT_4O} + + with pytest.raises(ValueError, match="Unsupported handoff type"): + AssistantAgent( + name="test_agent", + model_client=model_client, + handoffs=[123], # type: ignore + ) + + +class TestAssistantAgentStateManagement: + """Test suite for AssistantAgent state management.""" + + @pytest.mark.asyncio + async def test_save_and_load_state(self) -> None: + """Test saving and loading agent state.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + # Mock model context state + mock_context = MagicMock() + mock_context.save_state = AsyncMock(return_value={"context": "state"}) + mock_context.load_state = AsyncMock() + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + model_context=mock_context, + ) + + # Test save state + state = await agent.save_state() + assert "llm_context" in state + + # Test load state + await agent.load_state(state) + mock_context.load_state.assert_called_once() + + @pytest.mark.asyncio + async def test_on_reset(self) -> None: + """Test agent reset functionality.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + mock_context = MagicMock() + mock_context.clear = AsyncMock() + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + model_context=mock_context, + ) + + cancellation_token = CancellationToken() + await agent.on_reset(cancellation_token) + + mock_context.clear.assert_called_once() + + +class TestAssistantAgentProperties: + """Test suite for AssistantAgent properties.""" + + @pytest.mark.asyncio + async def test_produced_message_types_text_only(self) -> None: + """Test produced message types for text-only agent.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + ) + + message_types = agent.produced_message_types + assert TextMessage in message_types + + @pytest.mark.asyncio + async def test_produced_message_types_with_tools(self) -> None: + """Test produced message types for agent with tools.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": True, "vision": False, "family": ModelFamily.GPT_4O} + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + ) + + message_types = agent.produced_message_types + assert ToolCallSummaryMessage in message_types + + @pytest.mark.asyncio + async def test_produced_message_types_with_handoffs(self) -> None: + """Test produced message types for agent with handoffs.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": True, "vision": False, "family": ModelFamily.GPT_4O} + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + handoffs=["agent1"], + ) + + message_types = agent.produced_message_types + assert HandoffMessage in message_types + + @pytest.mark.asyncio + async def test_model_context_property(self) -> None: + """Test model_context property access.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + custom_context = BufferedChatCompletionContext(buffer_size=3) + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + model_context=custom_context, + ) + + assert agent.model_context == custom_context + + +class TestAssistantAgentErrorHandling: + """Test suite for error handling scenarios.""" + + @pytest.mark.asyncio + async def test_invalid_json_in_tool_arguments(self) -> None: + """Test handling of invalid JSON in tool arguments.""" + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments="invalid json", name="mock_tool_function")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + ) + + result = await agent.run(task="Execute tool with invalid JSON") + + # Should handle JSON parsing error + assert isinstance(result.messages[-1], ToolCallSummaryMessage) + + +class TestAssistantAgentMemoryIntegration: + """Test suite for AssistantAgent memory integration. + + Tests the integration between AssistantAgent and memory components, including: + - Memory initialization + - Context updates + - Query operations + - Memory persistence + """ + + @pytest.mark.asyncio + async def test_memory_updates_context(self) -> None: + """Test that memory properly updates model context. + + Verifies that: + 1. Memory contents are added to context + 2. Context updates trigger appropriate events + 3. Memory query results are properly handled + """ + # Setup test memory with initial content + memory = MockMemory(contents=["Previous conversation about topic A"]) + + # Configure model client with expected response + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="stop", + content="Response incorporating memory content", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + # Create agent with memory + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + memory=[memory], + description="Agent with memory integration", + ) + + # Track memory events during execution + memory_events: List[MemoryQueryEvent] = [] + + async def event_handler(event: MemoryQueryEvent) -> None: + """Handle memory query events. + + Args: + event: Memory query event to process + """ + memory_events.append(event) + + # Create a handler function to capture memory events + async def handle_memory_events(result: Any) -> None: + messages: List[BaseChatMessage] = result.messages if hasattr(result, "messages") else [] + for msg in messages: + if isinstance(msg, MemoryQueryEvent): + await event_handler(msg) + + # Run agent + result = await agent.run(task="Respond using memory context") + + # Process the events + await handle_memory_events(result) + + # Verify memory integration + assert len(memory_events) > 0, "No memory events were generated" + assert isinstance(result.messages[-1], TextMessage) + assert "Response incorporating memory content" in result.messages[-1].content + + @pytest.mark.asyncio + async def test_memory_persistence(self) -> None: + """Test memory persistence across multiple sessions. + + Verifies: + 1. Memory content persists between sessions + 2. Memory updates are preserved + 3. Context is properly restored + 4. Memory query events are generated correctly + """ + # Create memory with initial content + memory = MockMemory(contents=["Initial memory"]) + + # Create model client + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="stop", + content="Response using memory", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + CreateResult( + finish_reason="stop", + content="Response with updated memory", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": False, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + # Create agent with memory + agent = AssistantAgent(name="memory_test_agent", model_client=model_client, memory=[memory]) + + # First session + result1 = await agent.run(task="First task") + state = await agent.save_state() + + # Add new memory content + await memory.add(MemoryContent(content="New memory", mime_type="text/plain")) + + # Create new agent and restore state + new_agent = AssistantAgent(name="memory_test_agent", model_client=model_client, memory=[memory]) + await new_agent.load_state(state) + + # Second session + result2 = await new_agent.run(task="Second task") + + # Verify memory persistence + assert isinstance(result1.messages[-1], TextMessage) + assert isinstance(result2.messages[-1], TextMessage) + assert result1.messages[-1].content == "Response using memory" + assert result2.messages[-1].content == "Response with updated memory" + + # Verify memory events + memory_events = [msg for msg in result2.messages if isinstance(msg, MemoryQueryEvent)] + assert len(memory_events) > 0 + assert any("New memory" in str(event.content) for event in memory_events) + + +class TestAssistantAgentSystemMessage: + """Test suite for system message functionality.""" + + @pytest.mark.asyncio + async def test_system_message_none(self) -> None: + """Test agent with system_message=None.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + system_message=None, + ) + + assert agent._system_messages == [] # type: ignore[reportPrivateUsage] + + @pytest.mark.asyncio + async def test_custom_system_message(self) -> None: + """Test agent with custom system message.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + custom_message = "You are a specialized assistant." + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + system_message=custom_message, + ) + + assert len(agent._system_messages) == 1 # type: ignore[reportPrivateUsage] + assert agent._system_messages[0].content == custom_message # type: ignore[reportPrivateUsage] + + +class TestAssistantAgentModelCompatibility: + """Test suite for model compatibility functionality.""" + + @pytest.mark.asyncio + async def test_vision_compatibility(self) -> None: + """Test vision model compatibility.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": True, "family": ModelFamily.GPT_4O} + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + ) + + # Test _get_compatible_context with vision model + from autogen_core.models import LLMMessage + + messages: List[LLMMessage] = [SystemMessage(content="Test")] + compatible_messages = agent._get_compatible_context(model_client, messages) # type: ignore[reportPrivateUsage] + + # Should return original messages for vision models + assert compatible_messages == messages + + +class TestAssistantAgentComponentSerialization: + """Test suite for component serialization functionality.""" + + @pytest.mark.asyncio + async def test_to_config_basic_agent(self) -> None: + """Test _to_config method with basic agent configuration.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + model_client.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_client"}) + ) + + mock_context = MagicMock() + mock_context.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_context"}) + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + description="Test description", + system_message="Test system message", + model_context=mock_context, + metadata={"key": "value"}, + ) + + config = agent._to_config() # type: ignore[reportPrivateUsage] + + assert config.name == "test_agent" + assert config.description == "Test description" + assert config.system_message == "Test system message" + assert config.model_client_stream is False + assert config.reflect_on_tool_use is False + assert config.max_tool_iterations == 1 + assert config.metadata == {"key": "value"} + model_client.dump_component.assert_called_once() + mock_context.dump_component.assert_called_once() + + @pytest.mark.asyncio + async def test_to_config_agent_with_handoffs(self) -> None: + """Test _to_config method with agent having handoffs.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": True, "vision": False, "family": ModelFamily.GPT_4O} + model_client.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_client"}) + ) + + mock_context = MagicMock() + mock_context.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_context"}) + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + handoffs=["agent1", Handoff(target="agent2")], + model_context=mock_context, + ) + + config = agent._to_config() # type: ignore[reportPrivateUsage] + + assert config.handoffs is not None + assert len(config.handoffs) == 2 + handoff_targets: List[str] = [h.target if hasattr(h, "target") else str(h) for h in config.handoffs] # type: ignore[reportUnknownMemberType, reportAttributeAccessIssue] + assert "agent1" in handoff_targets + assert "agent2" in handoff_targets + + @pytest.mark.asyncio + async def test_to_config_agent_with_memory(self) -> None: + """Test _to_config method with agent having memory modules.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + model_client.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_client"}) + ) + + mock_context = MagicMock() + mock_context.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_context"}) + ) + + mock_memory = MockMemory() + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + memory=[mock_memory], + model_context=mock_context, + ) + + config = agent._to_config() # type: ignore[reportPrivateUsage] + + assert config.memory is not None + assert len(config.memory) == 1 + assert config.memory[0].provider == "test" + assert config.memory[0].config == {"type": "mock_memory"} + + @pytest.mark.asyncio + async def test_to_config_agent_with_workbench(self) -> None: + """Test _to_config method with agent having workbench.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": True, "vision": False, "family": ModelFamily.GPT_4O} + model_client.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_client"}) + ) + + mock_context = MagicMock() + mock_context.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_context"}) + ) + + mock_workbench = MagicMock() + mock_workbench.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_workbench"}) + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + model_context=mock_context, + ) + + # Replace the workbench with our mock + agent._workbench = [mock_workbench] # type: ignore[reportPrivateUsage] + + config = agent._to_config() # type: ignore[reportPrivateUsage] + + assert config.workbench is not None + assert len(config.workbench) == 1 + mock_workbench.dump_component.assert_called_once() + + @pytest.mark.asyncio + async def test_to_config_agent_with_structured_output(self) -> None: + """Test _to_config method with agent having structured output.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + model_client.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_client"}) + ) + + mock_context = MagicMock() + mock_context.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_context"}) + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + output_content_type=StructuredOutput, + model_context=mock_context, + ) + + config = agent._to_config() # type: ignore[reportPrivateUsage] + + assert config.structured_message_factory is not None + assert config.reflect_on_tool_use is True # Should be True with structured output + + @pytest.mark.asyncio + async def test_to_config_system_message_none(self) -> None: + """Test _to_config method with system_message=None.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + model_client.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_client"}) + ) + + mock_context = MagicMock() + mock_context.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_context"}) + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + system_message=None, + model_context=mock_context, + ) + + config = agent._to_config() # type: ignore[reportPrivateUsage] + + assert config.system_message is None + + @pytest.mark.asyncio + async def test_from_config_basic_agent(self) -> None: + """Test _from_config method with basic agent configuration.""" + mock_model_client = MagicMock() + mock_model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + with patch("autogen_core.models.ChatCompletionClient.load_component", return_value=mock_model_client): + config = AssistantAgentConfig( + name="test_agent", + model_client=ComponentModel(provider="test", config={"type": "mock_client"}), + description="Test description", + system_message="Test system", + model_client_stream=True, + reflect_on_tool_use=False, + tool_call_summary_format="{tool_name}: {result}", + metadata={"test": "value"}, + ) + + agent = AssistantAgent._from_config(config) # type: ignore[reportPrivateUsage] + + assert agent.name == "test_agent" + assert agent.description == "Test description" + assert agent._model_client_stream is True # type: ignore[reportPrivateUsage] + assert agent._reflect_on_tool_use is False # type: ignore[reportPrivateUsage] + assert agent._tool_call_summary_format == "{tool_name}: {result}" # type: ignore[reportPrivateUsage] + assert agent._metadata == {"test": "value"} # type: ignore[reportPrivateUsage] + + @pytest.mark.asyncio + async def test_from_config_with_structured_output(self) -> None: + """Test _from_config method with structured output configuration.""" + mock_model_client = MagicMock() + mock_model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + mock_structured_factory = MagicMock() + mock_structured_factory.format_string = "Test format" + mock_structured_factory.ContentModel = StructuredOutput + + with ( + patch("autogen_core.models.ChatCompletionClient.load_component", return_value=mock_model_client), + patch( + "autogen_agentchat.messages.StructuredMessageFactory.load_component", + return_value=mock_structured_factory, + ), + ): + config = AssistantAgentConfig( + name="test_agent", + model_client=ComponentModel(provider="test", config={"type": "mock_client"}), + description="Test description", + reflect_on_tool_use=True, + tool_call_summary_format="{result}", + structured_message_factory=ComponentModel(provider="test", config={"type": "mock_factory"}), + ) + + agent = AssistantAgent._from_config(config) # type: ignore[reportPrivateUsage] + + assert agent._reflect_on_tool_use is True # type: ignore[reportPrivateUsage] + assert agent._output_content_type == StructuredOutput # type: ignore[reportPrivateUsage] + assert agent._output_content_type_format == "Test format" # type: ignore[reportPrivateUsage] + + @pytest.mark.asyncio + async def test_from_config_with_workbench_and_memory(self) -> None: + """Test _from_config method with workbench and memory.""" + mock_model_client = MagicMock() + mock_model_client.model_info = {"function_calling": True, "vision": False, "family": ModelFamily.GPT_4O} + + mock_workbench = MagicMock() + mock_memory = MockMemory() + mock_context = MagicMock() + + with ( + patch("autogen_core.models.ChatCompletionClient.load_component", return_value=mock_model_client), + patch("autogen_core.tools.Workbench.load_component", return_value=mock_workbench), + patch("autogen_core.memory.Memory.load_component", return_value=mock_memory), + patch("autogen_core.model_context.ChatCompletionContext.load_component", return_value=mock_context), + ): + config = AssistantAgentConfig( + name="test_agent", + model_client=ComponentModel(provider="test", config={"type": "mock_client"}), + description="Test description", + workbench=[ComponentModel(provider="test", config={"type": "mock_workbench"})], + memory=[ComponentModel(provider="test", config={"type": "mock_memory"})], + model_context=ComponentModel(provider="test", config={"type": "mock_context"}), + reflect_on_tool_use=True, + tool_call_summary_format="{result}", + ) + + agent = AssistantAgent._from_config(config) # type: ignore[reportPrivateUsage] + + assert len(agent._workbench) == 1 # type: ignore[reportPrivateUsage] + assert agent._memory is not None # type: ignore[reportPrivateUsage] + assert len(agent._memory) == 1 # type: ignore[reportPrivateUsage] + assert agent._model_context == mock_context # type: ignore[reportPrivateUsage] + + @pytest.mark.asyncio + async def test_config_roundtrip_consistency(self) -> None: + """Test that converting to config and back preserves agent properties.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": True, "vision": False, "family": ModelFamily.GPT_4O} + model_client.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_client"}) + ) + + mock_context = MagicMock() + mock_context.dump_component = MagicMock( + return_value=ComponentModel(provider="test", config={"type": "mock_context"}) + ) + + original_agent = AssistantAgent( + name="test_agent", + model_client=model_client, + description="Test description", + system_message="Test system message", + model_client_stream=True, + reflect_on_tool_use=True, + max_tool_iterations=5, + tool_call_summary_format="{tool_name}: {result}", + handoffs=["agent1"], + model_context=mock_context, + metadata={"test": "value"}, + ) + + # Convert to config + config = original_agent._to_config() # type: ignore[reportPrivateUsage] + + # Verify config properties + assert config.name == "test_agent" + assert config.description == "Test description" + assert config.system_message == "Test system message" + assert config.model_client_stream is True + assert config.reflect_on_tool_use is True + assert config.max_tool_iterations == 5 + assert config.tool_call_summary_format == "{tool_name}: {result}" + assert config.metadata == {"test": "value"} + + +class TestAssistantAgentThoughtHandling: + """Test suite for thought handling functionality.""" + + @pytest.mark.asyncio + async def test_thought_event_yielded_from_model_result(self) -> None: + """Test that thought events are yielded when model result contains thoughts.""" + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="stop", + content="Final response", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + thought="This is my internal thought process", + ), + ], + model_info={ + "function_calling": False, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + ) + + messages: List[Any] = [] + async for message in agent.on_messages_stream( + [TextMessage(content="Test", source="user")], CancellationToken() + ): + messages.append(message) + + # Should have ThoughtEvent in the stream + thought_events = [msg for msg in messages if isinstance(msg, ThoughtEvent)] + assert len(thought_events) == 1 + assert thought_events[0].content == "This is my internal thought process" + assert thought_events[0].source == "test_agent" + + @pytest.mark.asyncio + async def test_thought_event_with_tool_calls(self) -> None: + """Test that thought events are yielded when tool calls have thoughts.""" + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="mock_tool_function")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + thought="I need to use this tool to help the user", + ), + CreateResult( + finish_reason="stop", + content="Tool execution completed", + usage=RequestUsage(prompt_tokens=15, completion_tokens=10), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + max_tool_iterations=1, + ) + + messages: List[Any] = [] + async for message in agent.on_messages_stream( + [TextMessage(content="Test", source="user")], CancellationToken() + ): + messages.append(message) + + # Should have ThoughtEvent in the stream + thought_events = [msg for msg in messages if isinstance(msg, ThoughtEvent)] + assert len(thought_events) == 1 + assert thought_events[0].content == "I need to use this tool to help the user" + assert thought_events[0].source == "test_agent" + + @pytest.mark.asyncio + async def test_thought_event_with_reflection(self) -> None: + """Test that thought events are yielded during reflection.""" + model_client = ReplayChatCompletionClient( + [ + # Initial tool call with thought + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="mock_tool_function")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + thought="Initial thought before tool call", + ), + # Reflection with thought + CreateResult( + finish_reason="stop", + content="Based on the tool result, here's my response", + usage=RequestUsage(prompt_tokens=15, completion_tokens=10), + cached=False, + thought="Reflection thought after tool execution", + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + reflect_on_tool_use=True, + model_client_stream=True, # Enable streaming + ) + + messages: List[Any] = [] + async for message in agent.on_messages_stream( + [TextMessage(content="Test", source="user")], CancellationToken() + ): + messages.append(message) + + # Should have two ThoughtEvents - one for initial call, one for reflection + thought_events = [msg for msg in messages if isinstance(msg, ThoughtEvent)] + assert len(thought_events) == 2 + + thought_contents = [event.content for event in thought_events] + assert "Initial thought before tool call" in thought_contents + assert "Reflection thought after tool execution" in thought_contents + + @pytest.mark.asyncio + async def test_thought_event_with_tool_call_loop(self) -> None: + """Test that thought events are yielded in tool call loops.""" + model_client = ReplayChatCompletionClient( + [ + # First tool call with thought + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "first"}), name="mock_tool_function")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + thought="First iteration thought", + ), + # Second tool call with thought + CreateResult( + finish_reason="function_calls", + content=[ + FunctionCall(id="2", arguments=json.dumps({"param": "second"}), name="mock_tool_function") + ], + usage=RequestUsage(prompt_tokens=12, completion_tokens=5), + cached=False, + thought="Second iteration thought", + ), + # Final response with thought + CreateResult( + finish_reason="stop", + content="Loop completed", + usage=RequestUsage(prompt_tokens=15, completion_tokens=10), + cached=False, + thought="Final completion thought", + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + max_tool_iterations=3, + ) + + messages: List[Any] = [] + async for message in agent.on_messages_stream( + [TextMessage(content="Test", source="user")], CancellationToken() + ): + messages.append(message) + + # Should have three ThoughtEvents - one for each iteration + thought_events = [msg for msg in messages if isinstance(msg, ThoughtEvent)] + assert len(thought_events) == 3 + + thought_contents = [event.content for event in thought_events] + assert "First iteration thought" in thought_contents + assert "Second iteration thought" in thought_contents + assert "Final completion thought" in thought_contents + + @pytest.mark.asyncio + async def test_thought_event_with_handoff(self) -> None: + """Test that thought events are included in handoff context.""" + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[ + FunctionCall( + id="1", arguments=json.dumps({"target": "other_agent"}), name="transfer_to_other_agent" + ) + ], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + thought="I need to hand this off to another agent", + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + handoffs=["other_agent"], + max_tool_iterations=1, + ) + + result = await agent.run(task="Test handoff with thought") + + # Should have ThoughtEvent in inner messages + thought_events = [msg for msg in result.messages if isinstance(msg, ThoughtEvent)] + assert len(thought_events) == 1 + assert thought_events[0].content == "I need to hand this off to another agent" + + # Should have handoff message with thought in context + handoff_message = result.messages[-1] + assert isinstance(handoff_message, HandoffMessage) + assert len(handoff_message.context) == 1 + assert isinstance(handoff_message.context[0], AssistantMessage) + assert handoff_message.context[0].content == "I need to hand this off to another agent" + + @pytest.mark.asyncio + async def test_no_thought_event_when_no_thought(self) -> None: + """Test that no thought events are yielded when model result has no thoughts.""" + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="stop", + content="Simple response without thought", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + # No thought field + ), + ], + model_info={ + "function_calling": False, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + ) + + messages: List[Any] = [] + async for message in agent.on_messages_stream( + [TextMessage(content="Test", source="user")], CancellationToken() + ): + messages.append(message) + + # Should have no ThoughtEvents + thought_events = [msg for msg in messages if isinstance(msg, ThoughtEvent)] + assert len(thought_events) == 0 + + @pytest.mark.asyncio + async def test_thought_event_context_preservation(self) -> None: + """Test that thoughts are properly preserved in model context.""" + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="stop", + content="Response with thought", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + thought="Internal reasoning", + ), + ], + model_info={ + "function_calling": False, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + ) + + await agent.run(task="Test thought preservation") + + # Check that the model context contains the thought + messages = await agent.model_context.get_messages() + assistant_messages = [msg for msg in messages if isinstance(msg, AssistantMessage)] + assert len(assistant_messages) > 0 + + # The last assistant message should have the thought + last_assistant_msg = assistant_messages[-1] + # Fix line 2730 - properly check for thought attribute with type checking + if hasattr(last_assistant_msg, "thought"): + thought_content = cast(str, last_assistant_msg.thought) + assert thought_content == "Internal reasoning" + + +class TestAssistantAgentAdvancedScenarios: + """Test suite for advanced usage scenarios.""" + + @pytest.mark.asyncio + async def test_handoff_without_tool_calls(self) -> None: + """Test handoff without any tool calls.""" + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[ + FunctionCall(id="1", arguments=json.dumps({"target": "agent2"}), name="transfer_to_agent2") + ], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + handoffs=["agent2"], + ) + + result = await agent.run(task="Handoff to agent2") + + # Should return HandoffMessage + assert isinstance(result.messages[-1], HandoffMessage) + assert result.messages[-1].target == "agent2" + + @pytest.mark.asyncio + async def test_multiple_handoff_warning(self) -> None: + """Test warning for multiple handoffs.""" + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[ + FunctionCall(id="1", arguments=json.dumps({"target": "agent2"}), name="transfer_to_agent2"), + FunctionCall(id="2", arguments=json.dumps({"target": "agent3"}), name="transfer_to_agent3"), + ], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + handoffs=["agent2", "agent3"], + ) + + with pytest.warns(UserWarning, match="Multiple handoffs detected"): + result = await agent.run(task="Multiple handoffs") + + # Should only execute first handoff + assert isinstance(result.messages[-1], HandoffMessage) + assert result.messages[-1].target == "agent2" + + @pytest.mark.asyncio + async def test_structured_output_with_reflection(self) -> None: + """Test structured output with reflection enabled.""" + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="mock_tool_function")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + CreateResult( + finish_reason="stop", + content='{"content": "Structured response", "confidence": 0.95}', + usage=RequestUsage(prompt_tokens=15, completion_tokens=10), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + output_content_type=StructuredOutput, + reflect_on_tool_use=True, + ) + + result = await agent.run(task="Test structured output with reflection") + + # Should return StructuredMessage + from autogen_agentchat.messages import StructuredMessage + + final_message = result.messages[-1] + assert isinstance(final_message, StructuredMessage) + # Fix line 1710 - properly access structured content with explicit type annotation + structured_message: StructuredMessage[StructuredOutput] = cast( + StructuredMessage[StructuredOutput], final_message + ) + assert structured_message.content.content == "Structured response" + assert structured_message.content.confidence == 0.95 + + +class TestAssistantAgentAdvancedToolFeatures: + """Test suite for advanced tool features including custom formatters.""" + + @pytest.mark.asyncio + async def test_custom_tool_call_summary_formatter(self) -> None: + """Test custom tool call summary formatter functionality.""" + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[ + FunctionCall(id="1", arguments=json.dumps({"param": "success"}), name="mock_tool_function"), + FunctionCall(id="2", arguments=json.dumps({"param": "error"}), name="mock_tool_function"), + ], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + def custom_formatter(call: FunctionCall, result: FunctionExecutionResult) -> str: + if result.is_error: + return f"ERROR in {call.name}: {result.content} (args: {call.arguments})" + else: + return f"SUCCESS: {call.name} completed" + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + tool_call_summary_formatter=custom_formatter, + reflect_on_tool_use=False, + ) + + result = await agent.run(task="Test custom formatter") + + # Should return ToolCallSummaryMessage with custom formatting + final_message = result.messages[-1] + assert isinstance(final_message, ToolCallSummaryMessage) + # Fix line 1875 - properly access content with type checking + assert hasattr(final_message, "content"), "ToolCallSummaryMessage should have content attribute" + content = final_message.content + assert "SUCCESS: mock_tool_function completed" in content + assert "SUCCESS: mock_tool_function completed" in content # Both calls should be successful + + @pytest.mark.asyncio + async def test_custom_tool_call_summary_format_string(self) -> None: + """Test custom tool call summary format string.""" + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="mock_tool_function")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + tool_call_summary_format="Tool {tool_name} called with {arguments} -> {result}", + reflect_on_tool_use=False, + ) + + result = await agent.run(task="Test custom format string") + + # Should return ToolCallSummaryMessage with custom format + final_message = result.messages[-1] + assert isinstance(final_message, ToolCallSummaryMessage) + content = final_message.content + assert "Tool mock_tool_function called with" in content + assert "Tool executed with: test" in content + + @pytest.mark.asyncio + async def test_tool_call_summary_formatter_overrides_format_string(self) -> None: + """Test that tool_call_summary_formatter overrides format string.""" + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="mock_tool_function")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + def custom_formatter(call: FunctionCall, result: FunctionExecutionResult) -> str: + return f"CUSTOM: {call.name} -> {result.content}" + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + tool_call_summary_format="This should be ignored: {result}", + tool_call_summary_formatter=custom_formatter, + reflect_on_tool_use=False, + ) -@pytest.mark.asyncio -async def test_model_client_stream_with_tool_calls() -> None: - mock_client = ReplayChatCompletionClient( - [ - CreateResult( - content=[ - FunctionCall(id="1", name="_pass_function", arguments=r'{"input": "task"}'), - FunctionCall(id="3", name="_echo_function", arguments=r'{"input": "task"}'), - ], - finish_reason="function_calls", + result = await agent.run(task="Test formatter override") + + # Should use custom formatter, not format string + final_message = result.messages[-1] + assert isinstance(final_message, ToolCallSummaryMessage) + content = final_message.content + assert "CUSTOM: mock_tool_function" in content + assert "This should be ignored" not in content + + @pytest.mark.asyncio + async def test_output_content_type_format_string(self) -> None: + """Test structured output with custom format string.""" + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="stop", + content='{"content": "Test response", "confidence": 0.8}', + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": False, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + output_content_type=StructuredOutput, + output_content_type_format="Response: {content} (Confidence: {confidence})", + ) + + result = await agent.run(task="Test structured output format") + + # Should return StructuredMessage with custom format + final_message = result.messages[-1] + assert isinstance(final_message, StructuredMessage) + # Fix line 1880 - properly access structured content with explicit type annotation + structured_message: StructuredMessage[StructuredOutput] = cast( + StructuredMessage[StructuredOutput], final_message + ) + assert structured_message.content.content == "Test response" + assert structured_message.content.confidence == 0.8 + # The format string should be stored in the agent + assert hasattr(agent, "_output_content_type_format") + output_format = getattr(agent, "_output_content_type_format", None) + assert output_format == "Response: {content} (Confidence: {confidence})" + + @pytest.mark.asyncio + async def test_tool_call_error_handling_with_custom_formatter(self) -> None: + """Test error handling in tool calls with custom formatter.""" + + def error_tool(param: str) -> str: + raise ValueError(f"Tool error with param: {param}") + + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="error_tool")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + def error_formatter(call: FunctionCall, result: FunctionExecutionResult) -> str: + if result.is_error: + return f"ERROR in {call.name}: {result.content}" + else: + return f"SUCCESS: {result.content}" + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[error_tool], + tool_call_summary_formatter=error_formatter, + reflect_on_tool_use=False, + ) + + result = await agent.run(task="Test error handling") + + # Should return ToolCallSummaryMessage with error formatting + assert isinstance(result.messages[-1], ToolCallSummaryMessage) + content = result.messages[-1].content + assert "ERROR in error_tool" in content + + @pytest.mark.asyncio + async def test_multiple_tools_with_different_formats(self) -> None: + """Test multiple tool calls with different return formats.""" + + def json_tool(data: str) -> str: + return json.dumps({"result": data, "status": "success"}) + + def simple_tool(text: str) -> str: + return f"Processed: {text}" + + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[ + FunctionCall(id="1", arguments=json.dumps({"data": "json_data"}), name="json_tool"), + FunctionCall(id="2", arguments=json.dumps({"text": "simple_text"}), name="simple_tool"), + ], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + def smart_formatter(call: FunctionCall, result: FunctionExecutionResult) -> str: + try: + # Try to parse as JSON + parsed = json.loads(result.content) + return f"{call.name}: {parsed}" + except json.JSONDecodeError: + # Plain text + return f"{call.name}: {result.content}" + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[json_tool, simple_tool], + tool_call_summary_formatter=smart_formatter, + reflect_on_tool_use=False, + ) + + result = await agent.run(task="Test multiple tool formats") + + # Should handle both JSON and plain text tools + assert isinstance(result.messages[-1], ToolCallSummaryMessage) + content = result.messages[-1].content + assert "json_tool:" in content + assert "simple_tool:" in content + assert "Processed: simple_text" in content + + +class TestAssistantAgentCancellationToken: + """Test suite for cancellation token handling.""" + + @pytest.mark.asyncio + async def test_cancellation_during_model_inference(self) -> None: + """Test cancellation token during model inference.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + # Mock create method to check cancellation token + model_client.create = AsyncMock() + model_client.create.return_value = CreateResult( + finish_reason="stop", + content="Response", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + ) + + cancellation_token = CancellationToken() + result = await agent.on_messages([TextMessage(content="Test", source="user")], cancellation_token) + + # Verify cancellation token was passed to model client + model_client.create.assert_called_once() + call_args = model_client.create.call_args + assert call_args.kwargs["cancellation_token"] == cancellation_token + # Verify result is not None + assert result is not None + + @pytest.mark.asyncio + async def test_cancellation_during_streaming_inference(self) -> None: + """Test cancellation token during streaming model inference.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + # Mock create_stream method + async def mock_create_stream(*args: Any, **kwargs: Any) -> Any: + yield "chunk1" # First chunk + yield "chunk2" # Second chunk + yield CreateResult( + finish_reason="stop", + content="chunk1chunk2", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + + model_client.create_stream = mock_create_stream + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + model_client_stream=True, + ) + + cancellation_token = CancellationToken() + messages: List[Any] = [] + async for message in agent.on_messages_stream([TextMessage(content="Test", source="user")], cancellation_token): + messages.append(message) + + # Should have received streaming chunks and final response + chunk_events = [msg for msg in messages if isinstance(msg, ModelClientStreamingChunkEvent)] + assert len(chunk_events) == 2 + assert chunk_events[0].content == "chunk1" + assert chunk_events[1].content == "chunk2" + + @pytest.mark.asyncio + async def test_cancellation_during_tool_execution(self) -> None: + """Test cancellation token during tool execution.""" + + async def slow_tool(param: str) -> str: + await asyncio.sleep(0.1) # Simulate slow operation + return f"Slow result: {param}" + + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="slow_tool")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[slow_tool], + ) + + cancellation_token = CancellationToken() + result = await agent.on_messages([TextMessage(content="Test", source="user")], cancellation_token) + + # Tool should execute successfully with cancellation token + assert isinstance(result.chat_message, ToolCallSummaryMessage) + assert "Slow result: test" in result.chat_message.content + + @pytest.mark.asyncio + async def test_cancellation_during_workbench_tool_execution(self) -> None: + """Test cancellation token during workbench tool execution.""" + mock_workbench = MagicMock() + mock_workbench.list_tools = AsyncMock(return_value=[{"name": "test_tool", "description": "Test tool"}]) + + # Mock tool execution result + mock_result = MagicMock() + mock_result.to_text.return_value = "Workbench tool result" + mock_result.is_error = False + mock_workbench.call_tool = AsyncMock(return_value=mock_result) + + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="test_tool")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + workbench=[mock_workbench], + ) + + cancellation_token = CancellationToken() + result = await agent.on_messages([TextMessage(content="Test", source="user")], cancellation_token) + + # Verify cancellation token was passed to workbench + mock_workbench.call_tool.assert_called_once() + call_args = mock_workbench.call_tool.call_args + assert call_args.kwargs["cancellation_token"] == cancellation_token + # Verify result is not None + assert result is not None + + @pytest.mark.asyncio + async def test_cancellation_during_memory_operations(self) -> None: + """Test cancellation token during memory operations.""" + mock_memory = MagicMock() + mock_memory.update_context = AsyncMock(return_value=None) + + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + model_client.create = AsyncMock( + return_value=CreateResult( + finish_reason="stop", + content="Response", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + memory=[mock_memory], + ) + + cancellation_token = CancellationToken() + await agent.on_messages([TextMessage(content="Test", source="user")], cancellation_token) + + # Memory update_context should be called + mock_memory.update_context.assert_called_once() + + @pytest.mark.asyncio + async def test_reset_with_cancellation_token(self) -> None: + """Test agent reset with cancellation token.""" + mock_context = MagicMock() + mock_context.clear = AsyncMock() + + agent = AssistantAgent( + name="test_agent", + model_client=MagicMock(), + model_context=mock_context, + ) + + cancellation_token = CancellationToken() + await agent.on_reset(cancellation_token) + + # Context clear should be called + mock_context.clear.assert_called_once() + + +class TestAssistantAgentStreamingEdgeCases: + """Test suite for streaming edge cases and error scenarios.""" + + @pytest.mark.asyncio + async def test_streaming_with_empty_chunks(self) -> None: + """Test streaming with empty chunks.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + async def mock_create_stream(*args: Any, **kwargs: Any) -> Any: + yield "" # Empty chunk + yield "content" + yield "" # Another empty chunk + yield CreateResult( + finish_reason="stop", + content="content", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + + model_client.create_stream = mock_create_stream + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + model_client_stream=True, + ) + + messages: List[Any] = [] + async for message in agent.on_messages_stream( + [TextMessage(content="Test", source="user")], CancellationToken() + ): + messages.append(message) + + # Should handle empty chunks gracefully + chunk_events = [msg for msg in messages if isinstance(msg, ModelClientStreamingChunkEvent)] + assert len(chunk_events) == 3 # Including empty chunks + assert chunk_events[0].content == "" + assert chunk_events[1].content == "content" + assert chunk_events[2].content == "" + + @pytest.mark.asyncio + async def test_streaming_with_invalid_chunk_type(self) -> None: + """Test streaming with invalid chunk type raises error.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + async def mock_create_stream(*args: Any, **kwargs: Any) -> Any: + yield "valid_chunk" + yield 123 # Invalid chunk type + yield CreateResult( + finish_reason="stop", + content="content", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + + model_client.create_stream = mock_create_stream + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + model_client_stream=True, + ) + + with pytest.raises(RuntimeError, match="Invalid chunk type"): + async for _ in agent.on_messages_stream([TextMessage(content="Test", source="user")], CancellationToken()): + pass + + @pytest.mark.asyncio + async def test_streaming_without_final_result(self) -> None: + """Test streaming without final CreateResult raises error.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + async def mock_create_stream(*args: Any, **kwargs: Any) -> Any: + yield "chunk1" + yield "chunk2" + # No final CreateResult + + model_client.create_stream = mock_create_stream + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + model_client_stream=True, + ) + + with pytest.raises(RuntimeError, match="No final model result in streaming mode"): + async for _ in agent.on_messages_stream([TextMessage(content="Test", source="user")], CancellationToken()): + pass + + @pytest.mark.asyncio + async def test_streaming_with_tool_calls_and_reflection(self) -> None: + """Test streaming with tool calls followed by reflection.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": True, "vision": False, "family": ModelFamily.GPT_4O} + + call_count = 0 + + async def mock_create_stream(*args: Any, **kwargs: Any) -> Any: + nonlocal call_count + call_count += 1 + + if call_count == 1: + # First call: tool call + yield CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="mock_tool_function")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + else: + # Second call: reflection streaming + yield "Reflection " + yield "response " + yield "complete" + yield CreateResult( + finish_reason="stop", + content="Reflection response complete", + usage=RequestUsage(prompt_tokens=15, completion_tokens=10), + cached=False, + ) + + model_client.create_stream = mock_create_stream + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + tools=[mock_tool_function], + reflect_on_tool_use=True, + model_client_stream=True, + ) + + messages: List[Any] = [] + async for message in agent.on_messages_stream( + [TextMessage(content="Test", source="user")], CancellationToken() + ): + messages.append(message) + + # Should have tool call events, execution events, and streaming chunks for reflection + tool_call_events = [msg for msg in messages if isinstance(msg, ToolCallRequestEvent)] + tool_exec_events = [msg for msg in messages if isinstance(msg, ToolCallExecutionEvent)] + chunk_events = [msg for msg in messages if isinstance(msg, ModelClientStreamingChunkEvent)] + + assert len(tool_call_events) == 1 + assert len(tool_exec_events) == 1 + assert len(chunk_events) == 3 # Three reflection chunks + assert chunk_events[0].content == "Reflection " + assert chunk_events[1].content == "response " + assert chunk_events[2].content == "complete" + + @pytest.mark.asyncio + async def test_streaming_with_large_chunks(self) -> None: + """Test streaming with large chunks.""" + model_client = MagicMock() + model_client.model_info = {"function_calling": False, "vision": False, "family": ModelFamily.GPT_4O} + + large_chunk = "x" * 10000 # 10KB chunk + + async def mock_create_stream(*args: Any, **kwargs: Any) -> Any: + yield large_chunk + yield CreateResult( + finish_reason="stop", + content=large_chunk, usage=RequestUsage(prompt_tokens=10, completion_tokens=5), cached=False, + ) + + model_client.create_stream = mock_create_stream + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + model_client_stream=True, + ) + + messages: List[Any] = [] + async for message in agent.on_messages_stream( + [TextMessage(content="Test", source="user")], CancellationToken() + ): + messages.append(message) + + # Should handle large chunks + chunk_events = [msg for msg in messages if isinstance(msg, ModelClientStreamingChunkEvent)] + assert len(chunk_events) == 1 + assert len(chunk_events[0].content) == 10000 + + +class TestAssistantAgentWorkbenchIntegration: + """Test suite for comprehensive workbench testing.""" + + @pytest.mark.asyncio + async def test_multiple_workbenches(self) -> None: + """Test agent with multiple workbenches.""" + mock_workbench1 = MagicMock() + mock_workbench1.list_tools = AsyncMock(return_value=[{"name": "tool1", "description": "Tool from workbench 1"}]) + mock_result1 = MagicMock() + mock_result1.to_text.return_value = "Result from workbench 1" + mock_result1.is_error = False + mock_workbench1.call_tool = AsyncMock(return_value=mock_result1) + + mock_workbench2 = MagicMock() + mock_workbench2.list_tools = AsyncMock(return_value=[{"name": "tool2", "description": "Tool from workbench 2"}]) + mock_result2 = MagicMock() + mock_result2.to_text.return_value = "Result from workbench 2" + mock_result2.is_error = False + mock_workbench2.call_tool = AsyncMock(return_value=mock_result2) + + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[ + FunctionCall(id="1", arguments=json.dumps({"param": "test1"}), name="tool1"), + FunctionCall(id="2", arguments=json.dumps({"param": "test2"}), name="tool2"), + ], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + workbench=[mock_workbench1, mock_workbench2], + ) + + result = await agent.run(task="Test multiple workbenches") + + # Both workbenches should be called + mock_workbench1.call_tool.assert_called_once() + mock_workbench2.call_tool.assert_called_once() + + # Should return summary with both results + assert isinstance(result.messages[-1], ToolCallSummaryMessage) + content = result.messages[-1].content + assert "Result from workbench 1" in content + assert "Result from workbench 2" in content + + @pytest.mark.asyncio + async def test_workbench_tool_not_found(self) -> None: + """Test handling when tool is not found in any workbench.""" + mock_workbench = MagicMock() + mock_workbench.list_tools = AsyncMock( + return_value=[{"name": "available_tool", "description": "Available tool"}] + ) + + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="missing_tool")], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + workbench=[mock_workbench], + ) + + result = await agent.run(task="Test missing tool") + + # Should return error message for missing tool + assert isinstance(result.messages[-1], ToolCallSummaryMessage) + content = result.messages[-1].content + assert "tool 'missing_tool' not found" in content + + @pytest.mark.asyncio + async def test_workbench_concurrent_tool_execution(self) -> None: + """Test concurrent execution of multiple workbench tools.""" + mock_workbench = MagicMock() + mock_workbench.list_tools = AsyncMock( + return_value=[ + {"name": "concurrent_tool1", "description": "Concurrent tool 1"}, + {"name": "concurrent_tool2", "description": "Concurrent tool 2"}, + ] + ) + + call_order: List[str] = [] + + async def mock_call_tool(name: str, **kwargs: Any) -> Any: + call_order.append(f"start_{name}") + await asyncio.sleep(0.01) # Simulate work + call_order.append(f"end_{name}") + + mock_result = MagicMock() + mock_result.to_text.return_value = f"Result from {name}" + mock_result.is_error = False + return mock_result + + mock_workbench.call_tool = mock_call_tool + + model_client = ReplayChatCompletionClient( + [ + CreateResult( + finish_reason="function_calls", + content=[ + FunctionCall(id="1", arguments=json.dumps({"param": "test1"}), name="concurrent_tool1"), + FunctionCall(id="2", arguments=json.dumps({"param": "test2"}), name="concurrent_tool2"), + ], + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="test_agent", + model_client=model_client, + workbench=[mock_workbench], + ) + + result = await agent.run(task="Test concurrent execution") + + # Should execute both tools concurrently (both start before either ends) + assert "start_concurrent_tool1" in call_order + assert "start_concurrent_tool2" in call_order + + # Both results should be present + assert isinstance(result.messages[-1], ToolCallSummaryMessage) + content = result.messages[-1].content + assert "Result from concurrent_tool1" in content + assert "Result from concurrent_tool2" in content + + +class TestAssistantAgentComplexIntegration: + """Test suite for complex integration scenarios.""" + + @pytest.mark.asyncio + async def test_complete_workflow_with_all_features(self) -> None: + """Test agent with tools, handoffs, memory, streaming, and reflection.""" + # Setup memory + memory = MockMemory(["User prefers detailed explanations"]) + + # Setup model client with complex workflow + model_client = ReplayChatCompletionClient( + [ + # Initial tool call + CreateResult( + finish_reason="function_calls", + content=[ + FunctionCall(id="1", arguments=json.dumps({"param": "analysis"}), name="mock_tool_function") + ], + usage=RequestUsage(prompt_tokens=20, completion_tokens=10), + cached=False, + thought="I need to analyze this first", + ), + # Reflection result + CreateResult( + finish_reason="stop", + content="Based on the analysis, I can provide a detailed response. The user prefers comprehensive explanations.", + usage=RequestUsage(prompt_tokens=30, completion_tokens=15), + cached=False, + thought="I should be thorough based on user preference", + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="comprehensive_agent", + model_client=model_client, + tools=[mock_tool_function], + handoffs=["specialist_agent"], + memory=[memory], + reflect_on_tool_use=True, + model_client_stream=True, + tool_call_summary_format="Analysis: {result}", + metadata={"test": "comprehensive"}, + ) + + messages: List[Any] = [] + async for message in agent.on_messages_stream( + [TextMessage(content="Analyze this complex scenario", source="user")], CancellationToken() + ): + messages.append(message) + + # Should have all types of events + memory_events = [msg for msg in messages if isinstance(msg, MemoryQueryEvent)] + thought_events = [msg for msg in messages if isinstance(msg, ThoughtEvent)] + tool_events = [msg for msg in messages if isinstance(msg, ToolCallRequestEvent)] + execution_events = [msg for msg in messages if isinstance(msg, ToolCallExecutionEvent)] + chunk_events = [msg for msg in messages if isinstance(msg, ModelClientStreamingChunkEvent)] + + assert len(memory_events) > 0 + assert len(thought_events) == 2 # Initial and reflection thoughts + assert len(tool_events) == 1 + assert len(execution_events) == 1 + assert len(chunk_events) == 0 # No streaming chunks since we removed the string responses + + # Final response should be TextMessage from reflection + final_response = None + for msg in reversed(messages): + if isinstance(msg, Response): + final_response = msg + break + + assert final_response is not None + assert isinstance(final_response.chat_message, TextMessage) + assert "comprehensive explanations" in final_response.chat_message.content + + @pytest.mark.asyncio + async def test_error_recovery_in_complex_workflow(self) -> None: + """Test error recovery in complex workflow with multiple failures.""" + + def failing_tool(param: str) -> str: + if param == "fail": + raise ValueError("Tool failure") + return f"Success: {param}" + + model_client = ReplayChatCompletionClient( + [ + # Multiple tool calls, some failing + CreateResult( + finish_reason="function_calls", + content=[ + FunctionCall(id="1", arguments=json.dumps({"param": "success"}), name="failing_tool"), + FunctionCall(id="2", arguments=json.dumps({"param": "fail"}), name="failing_tool"), + FunctionCall(id="3", arguments=json.dumps({"param": "success2"}), name="failing_tool"), + ], + usage=RequestUsage(prompt_tokens=20, completion_tokens=10), + cached=False, + ), + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + def error_aware_formatter(call: FunctionCall, result: FunctionExecutionResult) -> str: + if result.is_error: + return f"⚠️ {call.name} failed: {result.content}" + else: + return f"✅ {call.name}: {result.content}" + + agent = AssistantAgent( + name="error_recovery_agent", + model_client=model_client, + tools=[failing_tool], + tool_call_summary_formatter=error_aware_formatter, + reflect_on_tool_use=False, + ) + + result = await agent.run(task="Test error recovery") + + # Should handle mixed success/failure gracefully + assert isinstance(result.messages[-1], ToolCallSummaryMessage) + content = result.messages[-1].content + assert "✅ failing_tool: Success: success" in content + assert "⚠️ failing_tool failed:" in content + assert "✅ failing_tool: Success: success2" in content + + @pytest.mark.asyncio + async def test_state_persistence_across_interactions(self) -> None: + """Test that agent state persists correctly across multiple interactions.""" + model_client = ReplayChatCompletionClient( + [ + # First interaction + CreateResult( + finish_reason="stop", + content="First response", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + # Second interaction + CreateResult( + finish_reason="stop", + content="Second response, remembering context", + usage=RequestUsage(prompt_tokens=15, completion_tokens=8), + cached=False, + ), + ], + model_info={ + "function_calling": False, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4O, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + name="stateful_agent", + model_client=model_client, + system_message="Remember previous conversations", + ) + + # First interaction + result1 = await agent.run(task="First task") + final_message_1 = result1.messages[-1] + assert isinstance(final_message_1, TextMessage) + assert final_message_1.content == "First response" + + # Save state + state = await agent.save_state() + assert "llm_context" in state + + # Second interaction + result2 = await agent.run(task="Second task, referring to first") + # Fix line 2730 - properly access content on TextMessage + final_message_2 = result2.messages[-1] + assert isinstance(final_message_2, TextMessage) + assert final_message_2.content == "Second response, remembering context" + + # Verify context contains both interactions + context_messages = await agent.model_context.get_messages() + user_messages = [ + msg for msg in context_messages if hasattr(msg, "source") and getattr(msg, "source", None) == "user" + ] + assert len(user_messages) == 2 + + +class TestAssistantAgentMessageContext: + """Test suite for message context handling in AssistantAgent. + + Tests various scenarios of message handling, context updates, and state management. + """ + + @pytest.mark.asyncio + async def test_add_messages_to_context(self) -> None: + """Test adding different message types to context. + + Verifies: + 1. Regular messages are added correctly + 2. Handoff messages with context are handled properly + 3. Message order is preserved + 4. Model messages are converted correctly + """ + # Setup test context + model_context = BufferedChatCompletionContext(buffer_size=10) + + # Create test messages + regular_msg = TextMessage(content="Regular message", source="user") + handoff_msg = HandoffMessage(content="Handoff message", source="agent1", target="agent2") + + # Add messages to context + await AssistantAgent._add_messages_to_context(model_context=model_context, messages=[regular_msg, handoff_msg]) # type: ignore[reportPrivateUsage] + + # Verify context contents + context_messages = await model_context.get_messages() + + # Should have: regular + handoff = 2 messages (now that handoff doesn't have context) + assert len(context_messages) == 2 + + # Verify message order and content - only the added messages should be present + assert isinstance(context_messages[0], UserMessage) + assert context_messages[0].content == "Regular message" + + assert isinstance(context_messages[1], UserMessage) + assert context_messages[1].content == "Handoff message" + + # No more assertions needed for context_messages since we already verified both + + @pytest.mark.asyncio + async def test_complex_model_context(self) -> None: + """Test complex model context management scenarios. + + Verifies: + 1. Large context handling + 2. Mixed message type handling + 3. Context size limits + 4. Message filtering + """ + # Setup test context with limited size + model_context = BufferedChatCompletionContext(buffer_size=5) + + # Create a mix of message types + messages: List[BaseChatMessage] = [ + TextMessage(content="First message", source="user"), + StructuredMessage[StructuredOutput]( + content=StructuredOutput(content="Structured data", confidence=0.9), source="agent" ), - "Example response 2 to task", + ToolCallSummaryMessage(content="Tool result", source="agent", tool_calls=[], results=[]), + HandoffMessage(content="Handoff", source="agent1", target="agent2"), ] - ) - mock_client._model_info["function_calling"] = True # pyright: ignore - agent = AssistantAgent( - "test_agent", - model_client=mock_client, - model_client_stream=True, - reflect_on_tool_use=True, - tools=[_pass_function, _echo_function], - ) - chunks: List[str] = [] - async for message in agent.run_stream(task="task"): - if isinstance(message, TaskResult): - assert message.messages[-1].content == "Example response 2 to task" - assert message.messages[1].content == [ - FunctionCall(id="1", name="_pass_function", arguments=r'{"input": "task"}'), - FunctionCall(id="3", name="_echo_function", arguments=r'{"input": "task"}'), - ] - assert message.messages[2].content == [ - FunctionExecutionResult(call_id="1", content="pass", is_error=False, name="_pass_function"), - FunctionExecutionResult(call_id="3", content="task", is_error=False, name="_echo_function"), - ] - elif isinstance(message, ModelClientStreamingChunkEvent): - chunks.append(message.content) - assert "".join(chunks) == "Example response 2 to task" + + # Add messages to context + await AssistantAgent._add_messages_to_context(model_context=model_context, messages=messages) # type: ignore[reportPrivateUsage] + + # Verify context management + context_messages = await model_context.get_messages() + + # Should respect buffer size limit + assert len(context_messages) <= 5 + + # Verify message conversion + for msg in context_messages: + assert isinstance(msg, (SystemMessage, UserMessage, AssistantMessage)) + + +class TestAnthropicIntegration: + """Test suite for Anthropic model API integration.""" + + def _get_anthropic_client(self) -> AnthropicChatCompletionClient: + """Create an Anthropic client for testing.""" + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not found in environment variables") + + return AnthropicChatCompletionClient( + model="claude-3-haiku-20240307", # Use haiku for faster/cheaper testing + api_key=api_key, + temperature=0.0, + ) + + @pytest.mark.asyncio + async def test_anthropic_tool_call_loop_max_iterations_10(self) -> None: + """Test Anthropic integration with tool call loop and max_tool_iterations=10.""" + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not found in environment variables") + + client = self._get_anthropic_client() + + agent = AssistantAgent( + name="anthropic_test_agent", + model_client=client, + tools=[mock_tool_function], + max_tool_iterations=10, + ) + + # Test with a task that might require tool calls + result = await agent.run( + task="Use the mock_tool_function to process the text 'hello world'. Then provide a summary." + ) + + # Verify that we got a result + assert result is not None + assert isinstance(result, TaskResult) + assert len(result.messages) > 0 + # Check that the last message is a non-tool call. + assert isinstance(result.messages[-1], TextMessage) + # Check that a tool call was made + tool_calls = [msg for msg in result.messages if isinstance(msg, ToolCallRequestEvent)] + assert len(tool_calls) > 0 + + # Check that usage was tracked + usage = client.total_usage() + assert usage.prompt_tokens > 0 + assert usage.completion_tokens > 0 + + @pytest.mark.asyncio + async def test_anthropic_tool_call_loop_max_iterations_1_with_reflection(self) -> None: + """Test Anthropic integration with max_tool_iterations=1 and reflect_on_tool_use=True.""" + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not found in environment variables") + + client = self._get_anthropic_client() + + agent = AssistantAgent( + name="anthropic_reflection_agent", + model_client=client, + tools=[mock_tool_function], + max_tool_iterations=1, + reflect_on_tool_use=True, + ) + + # Test with a task that might require tool calls but should be limited to 1 iteration + result = await agent.run( + task="Use the mock_tool_function to process the text 'test input' and then explain what happened." + ) + + # Verify that we got a result + assert result is not None + assert isinstance(result, TaskResult) + assert len(result.messages) > 0 + # Check that the last message is a reflection + assert isinstance(result.messages[-1], TextMessage) + # Check that a tool call was made + tool_calls = [msg for msg in result.messages if isinstance(msg, ToolCallRequestEvent)] + assert len(tool_calls) > 0 + + # Check that usage was tracked + usage = client.total_usage() + assert usage.prompt_tokens > 0 + assert usage.completion_tokens > 0 + + @pytest.mark.asyncio + async def test_anthropic_basic_text_response(self) -> None: + """Test basic Anthropic integration without tools.""" + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not found in environment variables") + + client = self._get_anthropic_client() + + agent = AssistantAgent( + name="anthropic_basic_agent", + model_client=client, + ) + + # Test with a simple task that doesn't require tools + result = await agent.run(task="What is 2 + 2? Just answer with the number.") + + # Verify that we got a result + assert result is not None + assert isinstance(result, TaskResult) + # Check that we got a text message with content + assert isinstance(result.messages[-1], TextMessage) + assert "4" in result.messages[-1].content + + # Check that usage was tracked + usage = client.total_usage() + assert usage.prompt_tokens > 0 + assert usage.completion_tokens > 0 diff --git a/python/packages/autogen-agentchat/tests/test_code_executor_agent.py b/python/packages/autogen-agentchat/tests/test_code_executor_agent.py index 2ebf79feb4ae..1ed98c89e1d5 100644 --- a/python/packages/autogen-agentchat/tests/test_code_executor_agent.py +++ b/python/packages/autogen-agentchat/tests/test_code_executor_agent.py @@ -1,9 +1,20 @@ +import asyncio +from typing import List + import pytest from autogen_agentchat.agents import CodeExecutorAgent +from autogen_agentchat.agents._code_executor_agent import ApprovalFuncType, ApprovalRequest, ApprovalResponse from autogen_agentchat.base import Response -from autogen_agentchat.messages import TextMessage +from autogen_agentchat.messages import ( + CodeExecutionEvent, + CodeGenerationEvent, + TextMessage, +) from autogen_core import CancellationToken +from autogen_core.code_executor import CodeBlock +from autogen_core.models import ModelFamily, ModelInfo from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor +from autogen_ext.models.replay import ReplayChatCompletionClient @pytest.mark.asyncio @@ -34,6 +45,233 @@ async def test_basic_code_execution() -> None: assert response.chat_message.source == "code_executor" +@pytest.mark.asyncio +async def test_code_generation_and_execution_with_model_client() -> None: + """ + Tests the code generation, execution and reflection pipeline using a model client. + """ + + language = "python" + code = 'import math\n\nnumber = 42\nsquare_root = math.sqrt(number)\nprint("%0.3f" % (square_root,))' + + model_client = ReplayChatCompletionClient( + [f"Here is the code to calculate the square root of 42:\n```{language}\n{code}```".strip(), "TERMINATE"] + ) + + agent = CodeExecutorAgent( + name="code_executor_agent", code_executor=LocalCommandLineCodeExecutor(), model_client=model_client + ) + + messages = [ + TextMessage( + content="Generate python code to calculate the square root of 42", + source="assistant", + ) + ] + + code_generation_event: CodeGenerationEvent | None = None + code_execution_event: CodeExecutionEvent | None = None + response: Response | None = None + + async for message in agent.on_messages_stream(messages, CancellationToken()): + if isinstance(message, CodeGenerationEvent): + code_block = message.code_blocks[0] + assert code_block.code == code, "Code block does not match" + assert code_block.language == language, "Language does not match" + code_generation_event = message + elif isinstance(message, CodeExecutionEvent): + assert message.to_text().strip() == "6.481", f"Expected '6.481', got: {message.to_text().strip()}" + code_execution_event = message + elif isinstance(message, Response): + assert isinstance( + message.chat_message, TextMessage + ), f"Expected TextMessage, got: {type(message.chat_message)}" + assert ( + message.chat_message.source == "code_executor_agent" + ), f"Expected source 'code_executor_agent', got: {message.chat_message.source}" + response = message + else: + raise AssertionError(f"Unexpected message type: {type(message)}") + + assert code_generation_event is not None, "Code generation event was not received" + assert code_execution_event is not None, "Code execution event was not received" + assert response is not None, "Response was not received" + + +@pytest.mark.asyncio +async def test_no_code_response_with_model_client() -> None: + """ + Tests agent behavior when the model client responds with non-code content. + """ + + model_client = ReplayChatCompletionClient(["The capital of France is Paris.", "TERMINATE"]) + + agent = CodeExecutorAgent( + name="code_executor_agent", code_executor=LocalCommandLineCodeExecutor(), model_client=model_client + ) + + messages = [ + TextMessage( + content="What is the capital of France?", + source="assistant", + ) + ] + + response: Response | None = None + + async for message in agent.on_messages_stream(messages, CancellationToken()): + if isinstance(message, Response): + assert isinstance( + message.chat_message, TextMessage + ), f"Expected TextMessage, got: {type(message.chat_message)}" + assert ( + message.chat_message.source == "code_executor_agent" + ), f"Expected source 'code_executor_agent', got: {message.chat_message.source}" + assert ( + message.chat_message.content.strip() == "The capital of France is Paris." + ), f"Expected 'The capital of France is Paris.', got: {message.chat_message.content.strip()}" + response = message + else: + raise AssertionError(f"Unexpected message type: {type(message)}") + + assert response is not None, "Response was not received" + + +@pytest.mark.asyncio +async def test_self_debugging_loop() -> None: + """ + Tests self debugging loop when the model client responds with incorrect code. + """ + language = "python" + incorrect_code_block = """ +numbers = [10, 20, 30, 40, 50] +mean = sum(numbers) / len(numbers +print("The mean is:", mean) +""".strip() + incorrect_code_result = """ + mean = sum(numbers) / len(numbers + ^ +SyntaxError: '(' was never closed +""".strip() + correct_code_block = """ +numbers = [10, 20, 30, 40, 50] +mean = sum(numbers) / len(numbers) +print("The mean is:", mean) +""".strip() + correct_code_result = """ +The mean is: 30.0 +""".strip() + + model_client = ReplayChatCompletionClient( + [ + f""" +Here is the code to calculate the mean of 10, 20, 30, 40, 50 + +```{language} +{incorrect_code_block} +``` +""", + """{"retry": "true", "reason": "Retry 1: It is a test environment"}""", + f""" +Here is the updated code to calculate the mean of 10, 20, 30, 40, 50 + +```{language} +{correct_code_block} +```""", + "Final Response", + "TERMINATE", + ], + model_info=ModelInfo( + vision=False, + function_calling=False, + json_output=True, + family=ModelFamily.UNKNOWN, + structured_output=True, + ), + ) + + agent = CodeExecutorAgent( + name="code_executor_agent", + code_executor=LocalCommandLineCodeExecutor(), + model_client=model_client, + max_retries_on_error=1, + ) + + messages = [ + TextMessage( + content="Calculate the mean of 10, 20, 30, 40, 50.", + source="assistant", + ) + ] + + incorrect_code_generation_event: CodeGenerationEvent | None = None + correct_code_generation_event: CodeGenerationEvent | None = None + retry_decision_event: CodeGenerationEvent | None = None + incorrect_code_execution_event: CodeExecutionEvent | None = None + correct_code_execution_event: CodeExecutionEvent | None = None + response: Response | None = None + + message_id: int = 0 + async for message in agent.on_messages_stream(messages, CancellationToken()): + if isinstance(message, CodeGenerationEvent) and message_id == 0: + # Step 1: First code generation + code_block = message.code_blocks[0] + assert code_block.code.strip() == incorrect_code_block, "Incorrect code block does not match" + assert code_block.language == language, "Language does not match" + incorrect_code_generation_event = message + + elif isinstance(message, CodeExecutionEvent) and message_id == 1: + # Step 2: First code execution + assert ( + incorrect_code_result in message.to_text().strip() + ), f"Expected {incorrect_code_result} in execution result, got: {message.to_text().strip()}" + incorrect_code_execution_event = message + + elif isinstance(message, CodeGenerationEvent) and message_id == 2: + # Step 3: Retry generation with proposed correction + retry_response = "Attempt number: 1\nProposed correction: Retry 1: It is a test environment" + assert ( + message.to_text().strip() == retry_response + ), f"Expected {retry_response}, got: {message.to_text().strip()}" + retry_decision_event = message + + elif isinstance(message, CodeGenerationEvent) and message_id == 3: + # Step 4: Second retry code generation + code_block = message.code_blocks[0] + assert code_block.code.strip() == correct_code_block, "Correct code block does not match" + assert code_block.language == language, "Language does not match" + correct_code_generation_event = message + + elif isinstance(message, CodeExecutionEvent) and message_id == 4: + # Step 5: Second retry code execution + assert ( + message.to_text().strip() == correct_code_result + ), f"Expected {correct_code_result} in execution result, got: {message.to_text().strip()}" + correct_code_execution_event = message + + elif isinstance(message, Response) and message_id == 5: + # Step 6: Final response + assert isinstance( + message.chat_message, TextMessage + ), f"Expected TextMessage, got: {type(message.chat_message)}" + assert ( + message.chat_message.source == "code_executor_agent" + ), f"Expected source 'code_executor_agent', got: {message.chat_message.source}" + response = message + + else: + raise AssertionError(f"Unexpected message type: {type(message)}") + + message_id += 1 + + assert incorrect_code_generation_event is not None, "Incorrect code generation event was not received" + assert incorrect_code_execution_event is not None, "Incorrect code execution event was not received" + assert retry_decision_event is not None, "Retry decision event was not received" + assert correct_code_generation_event is not None, "Correct code generation event was not received" + assert correct_code_execution_event is not None, "Correct code execution event was not received" + assert response is not None, "Response was not received" + + @pytest.mark.asyncio async def test_code_execution_error() -> None: """Test basic code execution""" @@ -178,3 +416,219 @@ async def test_code_execution_agent_serialization() -> None: assert isinstance(deserialized_agent, CodeExecutorAgent) assert deserialized_agent.name == "code_executor" + + +@pytest.mark.asyncio +async def test_code_execution_agent_serialization_with_model_client() -> None: + """Test agent config serialization""" + + model_client = ReplayChatCompletionClient(["The capital of France is Paris.", "TERMINATE"]) + + agent = CodeExecutorAgent( + name="code_executor_agent", code_executor=LocalCommandLineCodeExecutor(), model_client=model_client + ) + + # Serialize and deserialize the agent + serialized_agent = agent.dump_component() + deserialized_agent = CodeExecutorAgent.load_component(serialized_agent) + + assert isinstance(deserialized_agent, CodeExecutorAgent) + assert deserialized_agent.name == "code_executor_agent" + assert deserialized_agent._model_client is not None # type: ignore + + +# Approval function test helpers +def approval_function_allow_all(request: ApprovalRequest) -> ApprovalResponse: + """Approval function that allows all code execution.""" + return ApprovalResponse(approved=True, reason="All code is approved") + + +def approval_function_deny_dangerous(request: ApprovalRequest) -> ApprovalResponse: + """Approval function that denies potentially dangerous code.""" + dangerous_keywords = ["rm ", "del ", "format", "delete", "DROP TABLE"] + + for keyword in dangerous_keywords: + if keyword in request.code: + return ApprovalResponse(approved=False, reason=f"Code contains potentially dangerous keyword: {keyword}") + + return ApprovalResponse(approved=True, reason="Code appears safe") + + +def approval_function_deny_all(request: ApprovalRequest) -> ApprovalResponse: + """Approval function that denies all code execution.""" + return ApprovalResponse(approved=False, reason="All code execution is denied") + + +# Async approval function test helpers +async def async_approval_function_allow_all(request: ApprovalRequest) -> ApprovalResponse: + """Async approval function that allows all code execution.""" + await asyncio.sleep(0.01) # Simulate async operation + return ApprovalResponse(approved=True, reason="All code is approved (async)") + + +async def async_approval_function_deny_dangerous(request: ApprovalRequest) -> ApprovalResponse: + """Async approval function that denies potentially dangerous code.""" + await asyncio.sleep(0.01) # Simulate async operation + dangerous_keywords = ["rm ", "del ", "format", "delete", "DROP TABLE"] + + for keyword in dangerous_keywords: + if keyword in request.code: + return ApprovalResponse( + approved=False, reason=f"Code contains potentially dangerous keyword: {keyword} (async)" + ) + + return ApprovalResponse(approved=True, reason="Code appears safe (async)") + + +async def async_approval_function_deny_all(request: ApprovalRequest) -> ApprovalResponse: + """Async approval function that denies all code execution.""" + await asyncio.sleep(0.01) # Simulate async operation + return ApprovalResponse(approved=False, reason="All code execution is denied (async)") + + +@pytest.mark.asyncio +async def test_approval_functionality_no_approval() -> None: + """Test that CodeExecutorAgent works without approval function (default behavior).""" + agent = CodeExecutorAgent("test_agent", LocalCommandLineCodeExecutor()) + + code_blocks = [CodeBlock(code="print('Hello World!')", language="python")] + result = await agent.execute_code_block(code_blocks, CancellationToken()) + + # Should execute successfully + assert result.exit_code == 0 + assert "Hello World!" in result.output + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "approval_func,code,language,expected_exit_code,expected_in_output", + [ + (approval_function_allow_all, "print('Approved code')", "python", 0, "Approved code"), + (approval_function_deny_dangerous, "print('Safe code')", "python", 0, "Safe code"), + (approval_function_deny_dangerous, "rm somefile.txt", "sh", 1, "dangerous keyword"), + (approval_function_deny_all, "print('This should be denied')", "python", 1, "All code execution is denied"), + ], +) +async def test_approval_functionality_sync( + approval_func: ApprovalFuncType, code: str, language: str, expected_exit_code: int, expected_in_output: str +) -> None: + """Test sync approval functionality with various approval functions and code samples.""" + agent = CodeExecutorAgent("test_agent", LocalCommandLineCodeExecutor(), approval_func=approval_func) + + code_blocks = [CodeBlock(code=code, language=language)] + result = await agent.execute_code_block(code_blocks, CancellationToken()) + + assert result.exit_code == expected_exit_code + assert expected_in_output in result.output + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +async def test_approval_functionality_context_passed(is_async: bool) -> None: + """Test that approval functions receive the correct context.""" + received_requests: List[ApprovalRequest] = [] + + if is_async: + + async def capture_context_async(request: ApprovalRequest) -> ApprovalResponse: + await asyncio.sleep(0.01) + received_requests.append(request) + return ApprovalResponse(approved=True, reason="Captured for testing (async)") + + agent = CodeExecutorAgent("test_agent", LocalCommandLineCodeExecutor(), approval_func=capture_context_async) + else: + + def capture_context_sync(request: ApprovalRequest) -> ApprovalResponse: + received_requests.append(request) + return ApprovalResponse(approved=True, reason="Captured for testing") + + agent = CodeExecutorAgent("test_agent", LocalCommandLineCodeExecutor(), approval_func=capture_context_sync) + + code_blocks = [CodeBlock(code="print('Test context')", language="python")] + await agent.execute_code_block(code_blocks, CancellationToken()) + + # Verify the approval function was called and received the correct data + assert len(received_requests) == 1 + request = received_requests[0] + assert isinstance(request, ApprovalRequest) + assert "print('Test context')" in request.code + assert "```python" in request.code + assert isinstance(request.context, list) + + +@pytest.mark.parametrize( + "approval_func", + [approval_function_allow_all, async_approval_function_allow_all], +) +def test_approval_functionality_serialization_fails(approval_func: ApprovalFuncType) -> None: + """Test that serialization fails when approval function is set.""" + agent = CodeExecutorAgent("test_agent", LocalCommandLineCodeExecutor(), approval_func=approval_func) + + # Should raise ValueError when trying to serialize + with pytest.raises(ValueError, match="Cannot serialize CodeExecutorAgent with approval_func set"): + agent.dump_component() + + +def test_approval_functionality_serialization_succeeds() -> None: + """Test that serialization succeeds when no approval function is set.""" + agent = CodeExecutorAgent("test_agent", LocalCommandLineCodeExecutor()) + + # Should serialize successfully + config = agent.dump_component() + assert config.config["name"] == "test_agent" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "approval_func,async_marker", + [ + (approval_function_deny_dangerous, ""), + (async_approval_function_deny_dangerous, "(async)"), + ], +) +async def test_approval_functionality_with_on_messages(approval_func: ApprovalFuncType, async_marker: str) -> None: + """Test approval functionality works with the on_messages interface.""" + agent = CodeExecutorAgent("test_agent", LocalCommandLineCodeExecutor(), approval_func=approval_func) + + # Test with safe code + safe_message = TextMessage(content="```python\nprint('Safe message')\n```", source="user") + response = await agent.on_messages([safe_message], CancellationToken()) + assert isinstance(response.chat_message, TextMessage) + assert "Safe message" in response.chat_message.content + + # Test with dangerous code + dangerous_message = TextMessage(content="```sh\nrm -rf /\n```", source="user") + response = await agent.on_messages([dangerous_message], CancellationToken()) + assert isinstance(response.chat_message, TextMessage) + assert "Code execution was not approved" in response.chat_message.content + if async_marker: + assert async_marker in response.chat_message.content + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "approval_func,code,language,expected_exit_code,expected_in_output", + [ + (async_approval_function_allow_all, "print('Approved async code')", "python", 0, "Approved async code"), + (async_approval_function_deny_dangerous, "print('Safe async code')", "python", 0, "Safe async code"), + (async_approval_function_deny_dangerous, "rm somefile.txt", "sh", 1, "dangerous keyword"), + ( + async_approval_function_deny_all, + "print('This should be denied async')", + "python", + 1, + "All code execution is denied (async)", + ), + ], +) +async def test_approval_functionality_async( + approval_func: ApprovalFuncType, code: str, language: str, expected_exit_code: int, expected_in_output: str +) -> None: + """Test async approval functionality with various approval functions and code samples.""" + agent = CodeExecutorAgent("test_agent", LocalCommandLineCodeExecutor(), approval_func=approval_func) + + code_blocks = [CodeBlock(code=code, language=language)] + result = await agent.execute_code_block(code_blocks, CancellationToken()) + + assert result.exit_code == expected_exit_code + assert expected_in_output in result.output diff --git a/python/packages/autogen-agentchat/tests/test_declarative_components.py b/python/packages/autogen-agentchat/tests/test_declarative_components.py index 4d7ba3f38bfb..69e5c4ae960c 100644 --- a/python/packages/autogen-agentchat/tests/test_declarative_components.py +++ b/python/packages/autogen-agentchat/tests/test_declarative_components.py @@ -14,8 +14,10 @@ from autogen_core.model_context import ( BufferedChatCompletionContext, HeadAndTailChatCompletionContext, + TokenLimitedChatCompletionContext, UnboundedChatCompletionContext, ) +from autogen_ext.models.openai import OpenAIChatCompletionClient @pytest.mark.asyncio @@ -104,6 +106,8 @@ async def test_chat_completion_context_declarative() -> None: unbounded_context = UnboundedChatCompletionContext() buffered_context = BufferedChatCompletionContext(buffer_size=5) head_tail_context = HeadAndTailChatCompletionContext(head_size=3, tail_size=2) + model_client = OpenAIChatCompletionClient(model="gpt-4o", api_key="test_key") + token_limited_context = TokenLimitedChatCompletionContext(model_client=model_client, token_limit=5) # Test serialization unbounded_config = unbounded_context.dump_component() @@ -118,6 +122,14 @@ async def test_chat_completion_context_declarative() -> None: assert head_tail_config.config["head_size"] == 3 assert head_tail_config.config["tail_size"] == 2 + token_limited_config = token_limited_context.dump_component() + assert token_limited_config.provider == "autogen_core.model_context.TokenLimitedChatCompletionContext" + assert token_limited_config.config["token_limit"] == 5 + assert ( + token_limited_config.config["model_client"]["provider"] + == "autogen_ext.models.openai.OpenAIChatCompletionClient" + ) + # Test deserialization loaded_unbounded = ComponentLoader.load_component(unbounded_config, UnboundedChatCompletionContext) assert isinstance(loaded_unbounded, UnboundedChatCompletionContext) @@ -129,3 +141,6 @@ async def test_chat_completion_context_declarative() -> None: loaded_head_tail = ComponentLoader.load_component(head_tail_config, HeadAndTailChatCompletionContext) assert isinstance(loaded_head_tail, HeadAndTailChatCompletionContext) + + loaded_token_limited = ComponentLoader.load_component(token_limited_config, TokenLimitedChatCompletionContext) + assert isinstance(loaded_token_limited, TokenLimitedChatCompletionContext) diff --git a/python/packages/autogen-agentchat/tests/test_events.py b/python/packages/autogen-agentchat/tests/test_events.py new file mode 100644 index 000000000000..b29f05caf22d --- /dev/null +++ b/python/packages/autogen-agentchat/tests/test_events.py @@ -0,0 +1,85 @@ +import json + +from autogen_agentchat.base import Response, TaskResult +from autogen_agentchat.messages import TextMessage +from autogen_agentchat.teams._group_chat._events import ( + GroupChatAgentResponse, + GroupChatMessage, + GroupChatStart, + GroupChatTeamResponse, +) + + +def test_group_chat_message_preserves_subclass_data() -> None: + """Test that GroupChatMessage preserves TextMessage subclass fields.""" + # Create a TextMessage with subclass-specific fields + text_msg = TextMessage( + content="Hello, world!", + source="TestAgent", + ) + + # Wrap in GroupChatMessage + group_msg = GroupChatMessage(message=text_msg) + + # Serialize and verify subclass fields are preserved + json_data = group_msg.model_dump_json() + parsed = json.loads(json_data) + + # The critical test: subclass fields should be preserved + assert "content" in parsed["message"], "TextMessage content field should be preserved" + assert "type" in parsed["message"], "TextMessage type field should be preserved" + assert parsed["message"]["content"] == "Hello, world!" + assert parsed["message"]["type"] == "TextMessage" + + +def test_group_chat_start_preserves_message_list_data() -> None: + """Test that GroupChatStart preserves subclass data in message lists.""" + text_msg1 = TextMessage(content="First message", source="Agent1") + text_msg2 = TextMessage(content="Second message", source="Agent2") + + group_start = GroupChatStart(messages=[text_msg1, text_msg2]) + + json_data = group_start.model_dump_json() + parsed = json.loads(json_data) + + # Check both messages preserve subclass data + assert "content" in parsed["messages"][0] + assert "content" in parsed["messages"][1] + assert parsed["messages"][0]["content"] == "First message" + assert parsed["messages"][1]["content"] == "Second message" + + +def test_group_chat_agent_response_preserves_dataclass_fields() -> None: + """Test that GroupChatAgentResponse preserves data in Response dataclass fields.""" + text_msg = TextMessage(content="Response message", source="ResponseAgent") + inner_text_msg = TextMessage(content="Inner message", source="InnerAgent") + response = Response(chat_message=text_msg, inner_messages=[inner_text_msg]) + + group_response = GroupChatAgentResponse(response=response, name="TestAgent") + + json_data = group_response.model_dump_json() + parsed = json.loads(json_data) + + # Verify dataclass field preserves subclass data + assert "content" in parsed["response"]["chat_message"] + assert "type" in parsed["response"]["chat_message"] + assert parsed["response"]["chat_message"]["content"] == "Response message" + inner_msgs = parsed["response"]["inner_messages"] + assert len(inner_msgs) == 1 + assert "content" in inner_msgs[0] + assert inner_msgs[0]["content"] == "Inner message" + + +def test_group_chat_team_response_preserves_nested_data() -> None: + """Test that GroupChatTeamResponse preserves deeply nested subclass data.""" + text_msg = TextMessage(content="Nested message", source="NestedAgent") + task_result = TaskResult(messages=[text_msg]) + + team_response = GroupChatTeamResponse(result=task_result, name="TestTeam") + + json_data = team_response.model_dump_json() + parsed = json.loads(json_data) + + # Verify deeply nested subclass data is preserved + assert "content" in parsed["result"]["messages"][0] + assert parsed["result"]["messages"][0]["content"] == "Nested message" diff --git a/python/packages/autogen-agentchat/tests/test_group_chat.py b/python/packages/autogen-agentchat/tests/test_group_chat.py index 5ab1605a72b7..3ded2e0c2e60 100644 --- a/python/packages/autogen-agentchat/tests/test_group_chat.py +++ b/python/packages/autogen-agentchat/tests/test_group_chat.py @@ -2,7 +2,7 @@ import json import logging import tempfile -from typing import AsyncGenerator, List, Sequence +from typing import Any, AsyncGenerator, Dict, List, Mapping, Sequence import pytest import pytest_asyncio @@ -12,14 +12,23 @@ BaseChatAgent, CodeExecutorAgent, ) -from autogen_agentchat.base import Handoff, Response, TaskResult -from autogen_agentchat.conditions import HandoffTermination, MaxMessageTermination, TextMentionTermination +from autogen_agentchat.base import Handoff, Response, TaskResult, TerminationCondition +from autogen_agentchat.conditions import ( + HandoffTermination, + MaxMessageTermination, + StopMessageTermination, + TextMentionTermination, +) from autogen_agentchat.messages import ( - AgentEvent, - ChatMessage, + BaseAgentEvent, + BaseChatMessage, HandoffMessage, + ModelClientStreamingChunkEvent, MultiModalMessage, + SelectorEvent, + SelectSpeakerEvent, StopMessage, + StructuredMessage, TextMessage, ToolCallExecutionEvent, ToolCallRequestEvent, @@ -31,6 +40,7 @@ from autogen_agentchat.teams._group_chat._swarm_group_chat import SwarmGroupChatManager from autogen_agentchat.ui import Console from autogen_core import AgentId, AgentRuntime, CancellationToken, FunctionCall, SingleThreadedAgentRuntime +from autogen_core.model_context import BufferedChatCompletionContext from autogen_core.models import ( AssistantMessage, CreateResult, @@ -44,7 +54,8 @@ from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.models.replay import ReplayChatCompletionClient -from utils import FileLogHandler +from pydantic import BaseModel +from utils import FileLogHandler, compare_messages, compare_task_results logger = logging.getLogger(EVENT_LOGGER_NAME) logger.setLevel(logging.DEBUG) @@ -58,14 +69,14 @@ def __init__(self, name: str, description: str) -> None: self._total_messages = 0 @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: return (TextMessage,) @property def total_messages(self) -> int: return self._total_messages - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: if len(messages) > 0: assert isinstance(messages[0], TextMessage) self._last_message = messages[0].content @@ -79,6 +90,16 @@ async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: async def on_reset(self, cancellation_token: CancellationToken) -> None: self._last_message = None + async def save_state(self) -> Mapping[str, Any]: + return { + "last_message": self._last_message, + "total_messages": self._total_messages, + } + + async def load_state(self, state: Mapping[str, Any]) -> None: + self._last_message = state.get("last_message") + self._total_messages = state.get("total_messages", 0) + class _FlakyAgent(BaseChatAgent): def __init__(self, name: str, description: str) -> None: @@ -87,20 +108,75 @@ def __init__(self, name: str, description: str) -> None: self._total_messages = 0 @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: return (TextMessage,) @property def total_messages(self) -> int: return self._total_messages - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: raise ValueError("I am a flaky agent...") async def on_reset(self, cancellation_token: CancellationToken) -> None: self._last_message = None +class _FlakyTermination(TerminationCondition): + def __init__(self, raise_on_count: int) -> None: + self._raise_on_count = raise_on_count + self._count = 0 + + @property + def terminated(self) -> bool: + """Check if the termination condition has been reached""" + return False + + async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None: + self._count += 1 + if self._count == self._raise_on_count: + raise ValueError("I am a flaky termination...") + return None + + async def reset(self) -> None: + pass + + +class _UnknownMessageType(BaseChatMessage): + content: str + + def to_model_message(self) -> UserMessage: + raise NotImplementedError("This message type is not supported.") + + def to_model_text(self) -> str: + raise NotImplementedError("This message type is not supported.") + + def to_text(self) -> str: + raise NotImplementedError("This message type is not supported.") + + def dump(self) -> Mapping[str, Any]: + return {} + + @classmethod + def load(cls, data: Mapping[str, Any]) -> "_UnknownMessageType": + return cls(**data) + + +class _UnknownMessageTypeAgent(BaseChatAgent): + def __init__(self, name: str, description: str) -> None: + super().__init__(name, description) + + @property + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: + return (_UnknownMessageType,) + + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: + return Response(chat_message=_UnknownMessageType(content="Unknown message type", source=self.name)) + + async def on_reset(self, cancellation_token: CancellationToken) -> None: + pass + + class _StopAgent(_EchoAgent): def __init__(self, name: str, description: str, *, stop_at: int = 1) -> None: super().__init__(name, description) @@ -108,10 +184,10 @@ def __init__(self, name: str, description: str, *, stop_at: int = 1) -> None: self._stop_at = stop_at @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: return (TextMessage, StopMessage) - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: self._count += 1 if self._count < self._stop_at: return await super().on_messages(messages, cancellation_token) @@ -122,6 +198,19 @@ def _pass_function(input: str) -> str: return "pass" +class _InputTask1(BaseModel): + task: str + data: List[str] + + +class _InputTask2(BaseModel): + task: str + data: str + + +TaskType = str | List[BaseChatMessage] | BaseChatMessage + + @pytest_asyncio.fixture(params=["single_threaded", "embedded"]) # type: ignore async def runtime(request: pytest.FixtureRequest) -> AsyncGenerator[AgentRuntime | None, None]: if request.param == "single_threaded": @@ -164,66 +253,295 @@ async def test_round_robin_group_chat(runtime: AgentRuntime | None) -> None: "Hello, world!", "TERMINATE", ] - # Normalize the messages to remove \r\n and any leading/trailing whitespace. - normalized_messages = [ - msg.content.replace("\r\n", "\n").rstrip("\n") if isinstance(msg.content, str) else msg.content - for msg in result.messages - ] - - # Assert that all expected messages are in the collected messages - assert normalized_messages == expected_messages + for i in range(len(expected_messages)): + produced_message = result.messages[i] + assert isinstance(produced_message, TextMessage) + content = produced_message.content.replace("\r\n", "\n").rstrip("\n") + assert content == expected_messages[i] assert result.stop_reason is not None and result.stop_reason == "Text 'TERMINATE' mentioned" - # Test streaming. + # Test streaming with default output_task_messages=True. model_client.reset() - index = 0 await team.reset() + streamed_messages: List[BaseAgentEvent | BaseChatMessage] = [] + final_stream_result: TaskResult | None = None async for message in team.run_stream( task="Write a program that prints 'Hello, world!'", ): if isinstance(message, TaskResult): - assert message == result + final_stream_result = message else: - assert message == result.messages[index] - index += 1 + streamed_messages.append(message) + assert final_stream_result is not None + assert compare_task_results(final_stream_result, result) + # Verify streamed messages match the complete result.messages + assert len(streamed_messages) == len(result.messages) + for streamed_msg, expected_msg in zip(streamed_messages, result.messages, strict=False): + assert compare_messages(streamed_msg, expected_msg) # Test message input. # Text message. model_client.reset() - index = 0 await team.reset() result_2 = await team.run( task=TextMessage(content="Write a program that prints 'Hello, world!'", source="user") ) - assert result == result_2 + assert compare_task_results(result, result_2) # Test multi-modal message. model_client.reset() - index = 0 await team.reset() - result_2 = await team.run( - task=MultiModalMessage(content=["Write a program that prints 'Hello, world!'"], source="user") + task = MultiModalMessage(content=["Write a program that prints 'Hello, world!'"], source="user") + result_2 = await team.run(task=task) + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result_2.messages[0], MultiModalMessage) + assert result.messages[0].content == task.content[0] + assert len(result.messages[1:]) == len(result_2.messages[1:]) + for i in range(1, len(result.messages)): + assert compare_messages(result.messages[i], result_2.messages[i]) + + +@pytest.mark.asyncio +async def test_round_robin_group_chat_output_task_messages_false(runtime: AgentRuntime | None) -> None: + model_client = ReplayChatCompletionClient( + [ + 'Here is the program\n ```python\nprint("Hello, world!")\n```', + "TERMINATE", + ], + ) + with tempfile.TemporaryDirectory() as temp_dir: + code_executor_agent = CodeExecutorAgent( + "code_executor", code_executor=LocalCommandLineCodeExecutor(work_dir=temp_dir) + ) + coding_assistant_agent = AssistantAgent( + "coding_assistant", + model_client=model_client, + ) + termination = TextMentionTermination("TERMINATE") + team = RoundRobinGroupChat( + participants=[coding_assistant_agent, code_executor_agent], + termination_condition=termination, + runtime=runtime, + ) + result = await team.run( + task="Write a program that prints 'Hello, world!'", + output_task_messages=False, + ) + expected_messages = [ + 'Here is the program\n ```python\nprint("Hello, world!")\n```', + "Hello, world!", + "TERMINATE", + ] + for i in range(len(expected_messages)): + produced_message = result.messages[i] + assert isinstance(produced_message, TextMessage) + content = produced_message.content.replace("\r\n", "\n").rstrip("\n") + assert content == expected_messages[i] + + assert result.stop_reason is not None and result.stop_reason == "Text 'TERMINATE' mentioned" + + # Test streaming with output_task_messages=False. + model_client.reset() + await team.reset() + streamed_messages: List[BaseAgentEvent | BaseChatMessage] = [] + final_stream_result: TaskResult | None = None + async for message in team.run_stream( + task="Write a program that prints 'Hello, world!'", + output_task_messages=False, + ): + if isinstance(message, TaskResult): + final_stream_result = message + else: + streamed_messages.append(message) + assert final_stream_result is not None + assert compare_task_results(final_stream_result, result) + # Verify streamed messages match the complete result.messages excluding the first task message + assert len(streamed_messages) == len(result.messages) # Exclude task message + for streamed_msg, expected_msg in zip(streamed_messages, result.messages, strict=False): + assert compare_messages(streamed_msg, expected_msg) + + # Test message input with output_task_messages=False. + # Text message. + model_client.reset() + await team.reset() + streamed_messages_2: List[BaseAgentEvent | BaseChatMessage] = [] + final_stream_result_2: TaskResult | None = None + async for message in team.run_stream( + task=TextMessage(content="Write a program that prints 'Hello, world!'", source="user"), + output_task_messages=False, + ): + if isinstance(message, TaskResult): + final_stream_result_2 = message + else: + streamed_messages_2.append(message) + assert final_stream_result_2 is not None + assert compare_task_results(final_stream_result_2, result) + # Verify streamed messages match the complete result.messages excluding the first task message + assert len(streamed_messages_2) == len(result.messages) + for streamed_msg, expected_msg in zip(streamed_messages_2, result.messages, strict=False): + assert compare_messages(streamed_msg, expected_msg) + + # Test multi-modal message with output_task_messages=False. + model_client.reset() + await team.reset() + task = MultiModalMessage(content=["Write a program that prints 'Hello, world!'"], source="user") + streamed_messages_3: List[BaseAgentEvent | BaseChatMessage] = [] + final_stream_result_3: TaskResult | None = None + async for message in team.run_stream(task=task, output_task_messages=False): + if isinstance(message, TaskResult): + final_stream_result_3 = message + else: + streamed_messages_3.append(message) + assert final_stream_result_3 is not None + # Verify streamed messages exclude the task message + assert len(streamed_messages_3) == len(final_stream_result_3.messages) + for streamed_msg, expected_msg in zip(streamed_messages_3, final_stream_result_3.messages, strict=False): + assert compare_messages(streamed_msg, expected_msg) + + +@pytest.mark.asyncio +async def test_round_robin_group_chat_with_team_event(runtime: AgentRuntime | None) -> None: + model_client = ReplayChatCompletionClient( + [ + 'Here is the program\n ```python\nprint("Hello, world!")\n```', + "TERMINATE", + ], + ) + with tempfile.TemporaryDirectory() as temp_dir: + code_executor_agent = CodeExecutorAgent( + "code_executor", code_executor=LocalCommandLineCodeExecutor(work_dir=temp_dir) + ) + coding_assistant_agent = AssistantAgent( + "coding_assistant", + model_client=model_client, + ) + termination = TextMentionTermination("TERMINATE") + team = RoundRobinGroupChat( + participants=[coding_assistant_agent, code_executor_agent], + termination_condition=termination, + runtime=runtime, + emit_team_events=True, + ) + result = await team.run( + task="Write a program that prints 'Hello, world!'", ) - assert result.messages[0].content == result_2.messages[0].content[0] - assert result.messages[1:] == result_2.messages[1:] + assert len(result.messages) == 7 + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result.messages[1], SelectSpeakerEvent) + assert isinstance(result.messages[2], TextMessage) + assert isinstance(result.messages[3], SelectSpeakerEvent) + assert isinstance(result.messages[4], TextMessage) + assert isinstance(result.messages[5], SelectSpeakerEvent) + assert isinstance(result.messages[6], TextMessage) + + # Test streaming with default output_task_messages=True. + model_client.reset() + await team.reset() + streamed_messages: List[BaseAgentEvent | BaseChatMessage] = [] + final_stream_result: TaskResult | None = None + async for message in team.run_stream( + task="Write a program that prints 'Hello, world!'", + ): + if isinstance(message, TaskResult): + final_stream_result = message + else: + streamed_messages.append(message) + assert final_stream_result is not None + assert compare_task_results(final_stream_result, result) + # Verify streamed messages match the complete result.messages + assert len(streamed_messages) == len(result.messages) + for streamed_msg, expected_msg in zip(streamed_messages, result.messages, strict=False): + assert compare_messages(streamed_msg, expected_msg) + + # Test multi-modal message. + model_client.reset() + await team.reset() + task = MultiModalMessage(content=["Write a program that prints 'Hello, world!'"], source="user") + result_2 = await team.run(task=task) + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result_2.messages[0], MultiModalMessage) + assert result.messages[0].content == task.content[0] + assert len(result.messages[1:]) == len(result_2.messages[1:]) + for i in range(1, len(result.messages)): + assert compare_messages(result.messages[i], result_2.messages[i]) @pytest.mark.asyncio -async def test_round_robin_group_chat_state(runtime: AgentRuntime | None) -> None: +async def test_round_robin_group_chat_unknown_task_message_type(runtime: AgentRuntime | None) -> None: + model_client = ReplayChatCompletionClient([]) + agent1 = AssistantAgent("agent1", model_client=model_client) + agent2 = AssistantAgent("agent2", model_client=model_client) + termination = TextMentionTermination("TERMINATE") + team1 = RoundRobinGroupChat( + participants=[agent1, agent2], + termination_condition=termination, + runtime=runtime, + custom_message_types=[StructuredMessage[_InputTask2]], + ) + with pytest.raises(ValueError, match=r"Message type .*StructuredMessage\[_InputTask1\].* is not registered"): + await team1.run( + task=StructuredMessage[_InputTask1]( + content=_InputTask1(task="Write a program that prints 'Hello, world!'", data=["a", "b", "c"]), + source="user", + ) + ) + + +@pytest.mark.asyncio +async def test_round_robin_group_chat_unknown_agent_message_type() -> None: + model_client = ReplayChatCompletionClient(["Hello"]) + agent1 = AssistantAgent("agent1", model_client=model_client) + agent2 = _UnknownMessageTypeAgent("agent2", "I am an unknown message type agent") + termination = TextMentionTermination("TERMINATE") + team1 = RoundRobinGroupChat(participants=[agent1, agent2], termination_condition=termination) + with pytest.raises(RuntimeError, match=".* Message type .*UnknownMessageType.* not registered"): + await team1.run(task=TextMessage(content="Write a program that prints 'Hello, world!'", source="user")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "task", + [ + "Write a program that prints 'Hello, world!'", + [TextMessage(content="Write a program that prints 'Hello, world!'", source="user")], + [MultiModalMessage(content=["Write a program that prints 'Hello, world!'"], source="user")], + [ + StructuredMessage[_InputTask1]( + content=_InputTask1(task="Write a program that prints 'Hello, world!'", data=["a", "b", "c"]), + source="user", + ), + StructuredMessage[_InputTask2]( + content=_InputTask2(task="Write a program that prints 'Hello, world!'", data="a"), source="user" + ), + ], + ], + ids=["text", "text_message", "multi_modal_message", "structured_message"], +) +async def test_round_robin_group_chat_state(task: TaskType, runtime: AgentRuntime | None) -> None: model_client = ReplayChatCompletionClient( ["No facts", "No plan", "print('Hello, world!')", "TERMINATE"], ) agent1 = AssistantAgent("agent1", model_client=model_client) agent2 = AssistantAgent("agent2", model_client=model_client) termination = TextMentionTermination("TERMINATE") - team1 = RoundRobinGroupChat(participants=[agent1, agent2], termination_condition=termination, runtime=runtime) - await team1.run(task="Write a program that prints 'Hello, world!'") + team1 = RoundRobinGroupChat( + participants=[agent1, agent2], + termination_condition=termination, + runtime=runtime, + custom_message_types=[StructuredMessage[_InputTask1], StructuredMessage[_InputTask2]], + ) + await team1.run(task=task) state = await team1.save_state() agent3 = AssistantAgent("agent1", model_client=model_client) agent4 = AssistantAgent("agent2", model_client=model_client) - team2 = RoundRobinGroupChat(participants=[agent3, agent4], termination_condition=termination, runtime=runtime) + team2 = RoundRobinGroupChat( + participants=[agent3, agent4], + termination_condition=termination, + runtime=runtime, + custom_message_types=[StructuredMessage[_InputTask1], StructuredMessage[_InputTask2]], + ) await team2.load_state(state) state2 = await team2.save_state() assert state == state2 @@ -260,7 +578,7 @@ async def test_round_robin_group_chat_with_tools(runtime: AgentRuntime | None) - "TERMINATE", ], model_info={ - "family": "gpt-4o", + "family": "gpt-4.1-nano", "function_calling": True, "json_output": True, "vision": True, @@ -294,24 +612,23 @@ async def test_round_robin_group_chat_with_tools(runtime: AgentRuntime | None) - # Test streaming. await tool_use_agent._model_context.clear() # pyright: ignore model_client.reset() - index = 0 + result_index = 0 # Include task message in result since output_task_messages=True by default await team.reset() async for message in team.run_stream( task="Write a program that prints 'Hello, world!'", ): if isinstance(message, TaskResult): - assert message == result + assert compare_task_results(message, result) else: - assert message == result.messages[index] - index += 1 + assert compare_messages(message, result.messages[result_index]) + result_index += 1 # Test Console. await tool_use_agent._model_context.clear() # pyright: ignore model_client.reset() - index = 0 await team.reset() result2 = await Console(team.run_stream(task="Write a program that prints 'Hello, world!'")) - assert result2 == result + assert compare_task_results(result2, result) @pytest.mark.asyncio @@ -349,10 +666,8 @@ async def test_round_robin_group_chat_with_resume_and_reset(runtime: AgentRuntim assert result.stop_reason is not None -# TODO: add runtime fixture for testing with custom runtime once the issue regarding -# hanging on exception is resolved. @pytest.mark.asyncio -async def test_round_robin_group_chat_with_exception_raised() -> None: +async def test_round_robin_group_chat_with_exception_raised_from_agent(runtime: AgentRuntime | None) -> None: agent_1 = _EchoAgent("agent_1", description="echo agent 1") agent_2 = _FlakyAgent("agent_2", description="echo agent 2") agent_3 = _EchoAgent("agent_3", description="echo agent 3") @@ -360,9 +675,29 @@ async def test_round_robin_group_chat_with_exception_raised() -> None: team = RoundRobinGroupChat( participants=[agent_1, agent_2, agent_3], termination_condition=termination, + runtime=runtime, ) - with pytest.raises(ValueError, match="I am a flaky agent..."): + with pytest.raises(RuntimeError, match="I am a flaky agent..."): + await team.run( + task="Write a program that prints 'Hello, world!'", + ) + + +@pytest.mark.asyncio +async def test_round_robin_group_chat_with_exception_raised_from_termination_condition( + runtime: AgentRuntime | None, +) -> None: + agent_1 = _EchoAgent("agent_1", description="echo agent 1") + agent_2 = _FlakyAgent("agent_2", description="echo agent 2") + agent_3 = _EchoAgent("agent_3", description="echo agent 3") + team = RoundRobinGroupChat( + participants=[agent_1, agent_2, agent_3], + termination_condition=_FlakyTermination(raise_on_count=1), + runtime=runtime, + ) + + with pytest.raises(Exception, match="I am a flaky termination..."): await team.run( task="Write a program that prints 'Hello, world!'", ) @@ -453,6 +788,7 @@ async def test_selector_group_chat(runtime: AgentRuntime | None) -> None: task="Write a program that prints 'Hello, world!'", ) assert len(result.messages) == 6 + assert isinstance(result.messages[0], TextMessage) assert result.messages[0].content == "Write a program that prints 'Hello, world!'" assert result.messages[1].source == "agent3" assert result.messages[2].source == "agent2" @@ -464,28 +800,158 @@ async def test_selector_group_chat(runtime: AgentRuntime | None) -> None: # Test streaming. model_client.reset() agent1._count = 0 # pyright: ignore - index = 0 + result_index = 0 # Include task message in result since output_task_messages=True by default await team.reset() async for message in team.run_stream( task="Write a program that prints 'Hello, world!'", ): if isinstance(message, TaskResult): - assert message == result + assert compare_task_results(message, result) else: - assert message == result.messages[index] - index += 1 + assert compare_messages(message, result.messages[result_index]) + result_index += 1 # Test Console. model_client.reset() agent1._count = 0 # pyright: ignore - index = 0 await team.reset() result2 = await Console(team.run_stream(task="Write a program that prints 'Hello, world!'")) - assert result2 == result + assert compare_task_results(result2, result) + + +@pytest.mark.asyncio +async def test_selector_group_chat_with_model_context(runtime: AgentRuntime | None) -> None: + buffered_context = BufferedChatCompletionContext(buffer_size=5) + await buffered_context.add_message(UserMessage(content="[User] Prefilled message", source="user")) + + selector_group_chat_model_client = ReplayChatCompletionClient( + ["agent2", "agent1", "agent1", "agent2", "agent1", "agent2", "agent1"] + ) + agent_one_model_client = ReplayChatCompletionClient( + ["[Agent One] First generation", "[Agent One] Second generation", "[Agent One] Third generation", "TERMINATE"] + ) + agent_two_model_client = ReplayChatCompletionClient( + ["[Agent Two] First generation", "[Agent Two] Second generation", "[Agent Two] Third generation"] + ) + + agent1 = AssistantAgent("agent1", model_client=agent_one_model_client, description="Assistant agent 1") + agent2 = AssistantAgent("agent2", model_client=agent_two_model_client, description="Assistant agent 2") + + termination = TextMentionTermination("TERMINATE") + team = SelectorGroupChat( + participants=[agent1, agent2], + model_client=selector_group_chat_model_client, + termination_condition=termination, + runtime=runtime, + emit_team_events=True, + allow_repeated_speaker=True, + model_context=buffered_context, + ) + await team.run( + task="[GroupChat] Task", + ) + + messages_to_check = [ + "user: [User] Prefilled message", + "user: [GroupChat] Task", + "agent2: [Agent Two] First generation", + "agent1: [Agent One] First generation", + "agent1: [Agent One] Second generation", + "agent2: [Agent Two] Second generation", + "agent1: [Agent One] Third generation", + "agent2: [Agent Two] Third generation", + ] + + create_calls: List[Dict[str, Any]] = selector_group_chat_model_client.create_calls + for idx, call in enumerate(create_calls): + messages = call["messages"] + prompt = messages[0].content + prompt_lines = prompt.split("\n") + chat_history = [value for value in messages_to_check[max(0, idx - 3) : idx + 2]] + assert all( + line.strip() in prompt_lines for line in chat_history + ), f"Expected all lines {chat_history} to be in prompt, but got {prompt_lines}" + + +@pytest.mark.asyncio +async def test_selector_group_chat_with_team_event(runtime: AgentRuntime | None) -> None: + model_client = ReplayChatCompletionClient( + ["agent3", "agent2", "agent1", "agent2", "agent1"], + ) + agent1 = _StopAgent("agent1", description="echo agent 1", stop_at=2) + agent2 = _EchoAgent("agent2", description="echo agent 2") + agent3 = _EchoAgent("agent3", description="echo agent 3") + termination = TextMentionTermination("TERMINATE") + team = SelectorGroupChat( + participants=[agent1, agent2, agent3], + model_client=model_client, + termination_condition=termination, + runtime=runtime, + emit_team_events=True, + ) + result = await team.run( + task="Write a program that prints 'Hello, world!'", + ) + assert len(result.messages) == 11 + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result.messages[1], SelectSpeakerEvent) + assert isinstance(result.messages[2], TextMessage) + assert isinstance(result.messages[3], SelectSpeakerEvent) + assert isinstance(result.messages[4], TextMessage) + assert isinstance(result.messages[5], SelectSpeakerEvent) + assert isinstance(result.messages[6], TextMessage) + assert isinstance(result.messages[7], SelectSpeakerEvent) + assert isinstance(result.messages[8], TextMessage) + assert isinstance(result.messages[9], SelectSpeakerEvent) + assert isinstance(result.messages[10], StopMessage) + assert result.messages[0].content == "Write a program that prints 'Hello, world!'" + assert result.messages[1].content == ["agent3"] + assert result.messages[2].source == "agent3" + assert result.messages[3].content == ["agent2"] + assert result.messages[4].source == "agent2" + assert result.messages[5].content == ["agent1"] + assert result.messages[6].source == "agent1" + assert result.messages[7].content == ["agent2"] + assert result.messages[8].source == "agent2" + assert result.messages[9].content == ["agent1"] + assert result.messages[10].source == "agent1" + assert result.stop_reason is not None and result.stop_reason == "Text 'TERMINATE' mentioned" + + # Test streaming. + model_client.reset() + agent1._count = 0 # pyright: ignore + result_index = 0 # Include task message in result since output_task_messages=True by default + await team.reset() + async for message in team.run_stream( + task="Write a program that prints 'Hello, world!'", + ): + if isinstance(message, TaskResult): + assert compare_task_results(message, result) + else: + assert compare_messages(message, result.messages[result_index]) + result_index += 1 @pytest.mark.asyncio -async def test_selector_group_chat_state(runtime: AgentRuntime | None) -> None: +@pytest.mark.parametrize( + "task", + [ + "Write a program that prints 'Hello, world!'", + [TextMessage(content="Write a program that prints 'Hello, world!'", source="user")], + [MultiModalMessage(content=["Write a program that prints 'Hello, world!'"], source="user")], + [ + StructuredMessage[_InputTask1]( + content=_InputTask1(task="Write a program that prints 'Hello, world!'", data=["a", "b", "c"]), + source="user", + ), + StructuredMessage[_InputTask2]( + content=_InputTask2(task="Write a program that prints 'Hello, world!'", data="a"), source="user" + ), + ], + ], + ids=["text", "text_message", "multi_modal_message", "structured_message"], +) +async def test_selector_group_chat_state(task: TaskType, runtime: AgentRuntime | None) -> None: model_client = ReplayChatCompletionClient( ["agent1", "No facts", "agent2", "No plan", "agent1", "print('Hello, world!')", "agent2", "TERMINATE"], ) @@ -497,14 +963,18 @@ async def test_selector_group_chat_state(runtime: AgentRuntime | None) -> None: termination_condition=termination, model_client=model_client, runtime=runtime, + custom_message_types=[StructuredMessage[_InputTask1], StructuredMessage[_InputTask2]], ) - await team1.run(task="Write a program that prints 'Hello, world!'") + await team1.run(task=task) state = await team1.save_state() agent3 = AssistantAgent("agent1", model_client=model_client) agent4 = AssistantAgent("agent2", model_client=model_client) team2 = SelectorGroupChat( - participants=[agent3, agent4], termination_condition=termination, model_client=model_client + participants=[agent3, agent4], + termination_condition=termination, + model_client=model_client, + custom_message_types=[StructuredMessage[_InputTask1], StructuredMessage[_InputTask2]], ) await team2.load_state(state) state2 = await team2.save_state() @@ -545,6 +1015,7 @@ async def test_selector_group_chat_two_speakers(runtime: AgentRuntime | None) -> task="Write a program that prints 'Hello, world!'", ) assert len(result.messages) == 5 + assert isinstance(result.messages[0], TextMessage) assert result.messages[0].content == "Write a program that prints 'Hello, world!'" assert result.messages[1].source == "agent2" assert result.messages[2].source == "agent1" @@ -555,22 +1026,21 @@ async def test_selector_group_chat_two_speakers(runtime: AgentRuntime | None) -> # Test streaming. model_client.reset() agent1._count = 0 # pyright: ignore - index = 0 + result_index = 0 # Include task message in result since output_task_messages=True by default await team.reset() async for message in team.run_stream(task="Write a program that prints 'Hello, world!'"): if isinstance(message, TaskResult): - assert message == result + assert compare_task_results(message, result) else: - assert message == result.messages[index] - index += 1 + assert compare_messages(message, result.messages[result_index]) + result_index += 1 # Test Console. model_client.reset() agent1._count = 0 # pyright: ignore - index = 0 await team.reset() result2 = await Console(team.run_stream(task="Write a program that prints 'Hello, world!'")) - assert result2 == result + assert compare_task_results(result2, result) @pytest.mark.asyncio @@ -594,6 +1064,7 @@ async def test_selector_group_chat_two_speakers_allow_repeated(runtime: AgentRun ) result = await team.run(task="Write a program that prints 'Hello, world!'") assert len(result.messages) == 4 + assert isinstance(result.messages[0], TextMessage) assert result.messages[0].content == "Write a program that prints 'Hello, world!'" assert result.messages[1].source == "agent2" assert result.messages[2].source == "agent2" @@ -602,21 +1073,20 @@ async def test_selector_group_chat_two_speakers_allow_repeated(runtime: AgentRun # Test streaming. model_client.reset() - index = 0 + result_index = 0 # Include task message in result since output_task_messages=True by default await team.reset() async for message in team.run_stream(task="Write a program that prints 'Hello, world!'"): if isinstance(message, TaskResult): - assert message == result + assert compare_task_results(message, result) else: - assert message == result.messages[index] - index += 1 + assert compare_messages(message, result.messages[result_index]) + result_index += 1 # Test Console. model_client.reset() - index = 0 await team.reset() result2 = await Console(team.run_stream(task="Write a program that prints 'Hello, world!'")) - assert result2 == result + assert compare_task_results(result2, result) @pytest.mark.asyncio @@ -635,6 +1105,7 @@ async def test_selector_group_chat_succcess_after_2_attempts(runtime: AgentRunti ) result = await team.run(task="Write a program that prints 'Hello, world!'") assert len(result.messages) == 2 + assert isinstance(result.messages[0], TextMessage) assert result.messages[0].content == "Write a program that prints 'Hello, world!'" assert result.messages[1].source == "agent2" @@ -659,6 +1130,7 @@ async def test_selector_group_chat_fall_back_to_first_after_3_attempts(runtime: ) result = await team.run(task="Write a program that prints 'Hello, world!'") assert len(result.messages) == 2 + assert isinstance(result.messages[0], TextMessage) assert result.messages[0].content == "Write a program that prints 'Hello, world!'" assert result.messages[1].source == "agent1" @@ -679,6 +1151,7 @@ async def test_selector_group_chat_fall_back_to_previous_after_3_attempts(runtim ) result = await team.run(task="Write a program that prints 'Hello, world!'") assert len(result.messages) == 3 + assert isinstance(result.messages[0], TextMessage) assert result.messages[0].content == "Write a program that prints 'Hello, world!'" assert result.messages[1].source == "agent2" assert result.messages[2].source == "agent2" @@ -692,7 +1165,7 @@ async def test_selector_group_chat_custom_selector(runtime: AgentRuntime | None) agent3 = _EchoAgent("agent3", description="echo agent 3") agent4 = _EchoAgent("agent4", description="echo agent 4") - def _select_agent(messages: Sequence[AgentEvent | ChatMessage]) -> str | None: + def _select_agent(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> str | None: if len(messages) == 0: return "agent1" elif messages[-1].source == "agent1": @@ -733,7 +1206,7 @@ async def test_selector_group_chat_custom_candidate_func(runtime: AgentRuntime | agent3 = _EchoAgent("agent3", description="echo agent 3") agent4 = _EchoAgent("agent4", description="echo agent 4") - def _candidate_func(messages: Sequence[AgentEvent | ChatMessage]) -> List[str]: + def _candidate_func(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str]: if len(messages) == 0: return ["agent1"] elif messages[-1].source == "agent1": @@ -772,10 +1245,10 @@ def __init__(self, name: str, description: str, next_agent: str) -> None: self._next_agent = next_agent @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: return (HandoffMessage,) - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: return Response( chat_message=HandoffMessage( content=f"Transferred to {self._next_agent}.", target=self._next_agent, source=self.name @@ -796,6 +1269,12 @@ async def test_swarm_handoff(runtime: AgentRuntime | None) -> None: team = Swarm([second_agent, first_agent, third_agent], termination_condition=termination, runtime=runtime) result = await team.run(task="task") assert len(result.messages) == 6 + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result.messages[1], HandoffMessage) + assert isinstance(result.messages[2], HandoffMessage) + assert isinstance(result.messages[3], HandoffMessage) + assert isinstance(result.messages[4], HandoffMessage) + assert isinstance(result.messages[5], HandoffMessage) assert result.messages[0].content == "task" assert result.messages[1].content == "Transferred to third_agent." assert result.messages[2].content == "Transferred to first_agent." @@ -808,15 +1287,15 @@ async def test_swarm_handoff(runtime: AgentRuntime | None) -> None: ) # Test streaming. - index = 0 + result_index = 0 # Include task message in result since output_task_messages=True by default await team.reset() stream = team.run_stream(task="task") async for message in stream: if isinstance(message, TaskResult): - assert message == result + assert compare_task_results(message, result) else: - assert message == result.messages[index] - index += 1 + assert compare_messages(message, result.messages[result_index]) + result_index += 1 # Test save and load. state = await team.save_state() @@ -839,6 +1318,119 @@ async def test_swarm_handoff(runtime: AgentRuntime | None) -> None: assert manager_1._current_speaker == manager_2._current_speaker # pyright: ignore +@pytest.mark.asyncio +async def test_swarm_handoff_with_team_events(runtime: AgentRuntime | None) -> None: + first_agent = _HandOffAgent("first_agent", description="first agent", next_agent="second_agent") + second_agent = _HandOffAgent("second_agent", description="second agent", next_agent="third_agent") + third_agent = _HandOffAgent("third_agent", description="third agent", next_agent="first_agent") + + termination = MaxMessageTermination(6) + team = Swarm( + [second_agent, first_agent, third_agent], + termination_condition=termination, + runtime=runtime, + emit_team_events=True, + ) + result = await team.run(task="task") + assert len(result.messages) == 11 + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result.messages[1], SelectSpeakerEvent) + assert isinstance(result.messages[2], HandoffMessage) + assert isinstance(result.messages[3], SelectSpeakerEvent) + assert isinstance(result.messages[4], HandoffMessage) + assert isinstance(result.messages[5], SelectSpeakerEvent) + assert isinstance(result.messages[6], HandoffMessage) + assert isinstance(result.messages[7], SelectSpeakerEvent) + assert isinstance(result.messages[8], HandoffMessage) + assert isinstance(result.messages[9], SelectSpeakerEvent) + assert isinstance(result.messages[10], HandoffMessage) + assert result.messages[0].content == "task" + assert result.messages[1].content == ["second_agent"] + assert result.messages[2].content == "Transferred to third_agent." + assert result.messages[3].content == ["third_agent"] + assert result.messages[4].content == "Transferred to first_agent." + assert result.messages[5].content == ["first_agent"] + assert result.messages[6].content == "Transferred to second_agent." + assert result.messages[7].content == ["second_agent"] + assert result.messages[8].content == "Transferred to third_agent." + assert result.messages[9].content == ["third_agent"] + assert result.messages[10].content == "Transferred to first_agent." + assert ( + result.stop_reason is not None + and result.stop_reason == "Maximum number of messages 6 reached, current message count: 6" + ) + + # Test streaming. + result_index = 0 # Include task message in result since output_task_messages=True by default + await team.reset() + stream = team.run_stream(task="task") + async for message in stream: + if isinstance(message, TaskResult): + assert compare_task_results(message, result) + else: + assert compare_messages(message, result.messages[result_index]) + result_index += 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "task", + [ + "Write a program that prints 'Hello, world!'", + [TextMessage(content="Write a program that prints 'Hello, world!'", source="user")], + [MultiModalMessage(content=["Write a program that prints 'Hello, world!'"], source="user")], + [ + StructuredMessage[_InputTask1]( + content=_InputTask1(task="Write a program that prints 'Hello, world!'", data=["a", "b", "c"]), + source="user", + ), + StructuredMessage[_InputTask2]( + content=_InputTask2(task="Write a program that prints 'Hello, world!'", data="a"), source="user" + ), + ], + ], + ids=["text", "text_message", "multi_modal_message", "structured_message"], +) +async def test_swarm_handoff_state(task: TaskType, runtime: AgentRuntime | None) -> None: + first_agent = _HandOffAgent("first_agent", description="first agent", next_agent="second_agent") + second_agent = _HandOffAgent("second_agent", description="second agent", next_agent="third_agent") + third_agent = _HandOffAgent("third_agent", description="third agent", next_agent="first_agent") + + termination = MaxMessageTermination(6) + team1 = Swarm( + [second_agent, first_agent, third_agent], + termination_condition=termination, + runtime=runtime, + custom_message_types=[StructuredMessage[_InputTask1], StructuredMessage[_InputTask2]], + ) + await team1.run(task=task) + state = await team1.save_state() + + first_agent2 = _HandOffAgent("first_agent", description="first agent", next_agent="second_agent") + second_agent2 = _HandOffAgent("second_agent", description="second agent", next_agent="third_agent") + third_agent2 = _HandOffAgent("third_agent", description="third agent", next_agent="first_agent") + team2 = Swarm( + [second_agent2, first_agent2, third_agent2], + termination_condition=termination, + runtime=runtime, + custom_message_types=[StructuredMessage[_InputTask1], StructuredMessage[_InputTask2]], + ) + await team2.load_state(state) + state2 = await team2.save_state() + assert state == state2 + + manager_1 = await team1._runtime.try_get_underlying_agent_instance( # pyright: ignore + AgentId(f"{team1._group_chat_manager_name}_{team1._team_id}", team1._team_id), # pyright: ignore + SwarmGroupChatManager, # pyright: ignore + ) + manager_2 = await team2._runtime.try_get_underlying_agent_instance( # pyright: ignore + AgentId(f"{team2._group_chat_manager_name}_{team2._team_id}", team2._team_id), # pyright: ignore + SwarmGroupChatManager, # pyright: ignore + ) + assert manager_1._message_thread == manager_2._message_thread # pyright: ignore + assert manager_1._current_speaker == manager_2._current_speaker # pyright: ignore + + @pytest.mark.asyncio async def test_swarm_handoff_using_tool_calls(runtime: AgentRuntime | None) -> None: model_client = ReplayChatCompletionClient( @@ -853,7 +1445,7 @@ async def test_swarm_handoff_using_tool_calls(runtime: AgentRuntime | None) -> N "TERMINATE", ], model_info={ - "family": "gpt-4o", + "family": "gpt-4.1-nano", "function_calling": True, "json_output": True, "vision": True, @@ -870,9 +1462,14 @@ async def test_swarm_handoff_using_tool_calls(runtime: AgentRuntime | None) -> N team = Swarm([agent1, agent2], termination_condition=termination, runtime=runtime) result = await team.run(task="task") assert len(result.messages) == 7 + assert isinstance(result.messages[0], TextMessage) assert result.messages[0].content == "task" assert isinstance(result.messages[1], ToolCallRequestEvent) assert isinstance(result.messages[2], ToolCallExecutionEvent) + assert isinstance(result.messages[3], HandoffMessage) + assert isinstance(result.messages[4], HandoffMessage) + assert isinstance(result.messages[5], TextMessage) + assert isinstance(result.messages[6], TextMessage) assert result.messages[3].content == "handoff to agent2" assert result.messages[4].content == "Transferred to agent1." assert result.messages[5].content == "Hello" @@ -882,23 +1479,22 @@ async def test_swarm_handoff_using_tool_calls(runtime: AgentRuntime | None) -> N # Test streaming. await agent1._model_context.clear() # pyright: ignore model_client.reset() - index = 0 + result_index = 0 # Include task message in result since output_task_messages=True by default await team.reset() stream = team.run_stream(task="task") async for message in stream: if isinstance(message, TaskResult): - assert message == result + assert compare_task_results(message, result) else: - assert message == result.messages[index] - index += 1 + assert compare_messages(message, result.messages[result_index]) + result_index += 1 # Test Console await agent1._model_context.clear() # pyright: ignore model_client.reset() - index = 0 await team.reset() result2 = await Console(team.run_stream(task="task")) - assert result2 == result + assert compare_task_results(result2, result) @pytest.mark.asyncio @@ -910,18 +1506,23 @@ async def test_swarm_pause_and_resume(runtime: AgentRuntime | None) -> None: team = Swarm([second_agent, first_agent, third_agent], max_turns=1, runtime=runtime) result = await team.run(task="task") assert len(result.messages) == 2 + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result.messages[1], HandoffMessage) assert result.messages[0].content == "task" assert result.messages[1].content == "Transferred to third_agent." # Resume with a new task. result = await team.run(task="new task") assert len(result.messages) == 2 + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result.messages[1], HandoffMessage) assert result.messages[0].content == "new task" assert result.messages[1].content == "Transferred to first_agent." # Resume with the same task. result = await team.run() assert len(result.messages) == 1 + assert isinstance(result.messages[0], HandoffMessage) assert result.messages[0].content == "Transferred to second_agent." @@ -943,7 +1544,7 @@ async def test_swarm_with_parallel_tool_calls(runtime: AgentRuntime | None) -> N "TERMINATE", ], model_info={ - "family": "gpt-4o", + "family": "gpt-4.1-nano", "function_calling": True, "json_output": True, "vision": True, @@ -987,17 +1588,22 @@ def tool2() -> str: team = Swarm([agent1, agent2], termination_condition=termination, runtime=runtime) result = await team.run(task="task") assert len(result.messages) == 6 - assert result.messages[0] == TextMessage(content="task", source="user") + assert compare_messages(result.messages[0], TextMessage(content="task", source="user")) assert isinstance(result.messages[1], ToolCallRequestEvent) assert isinstance(result.messages[2], ToolCallExecutionEvent) - assert result.messages[3] == HandoffMessage( - content="handoff to agent2", - target="agent2", - source="agent1", - context=expected_handoff_context, + assert compare_messages( + result.messages[3], + HandoffMessage( + content="handoff to agent2", + target="agent2", + source="agent1", + context=expected_handoff_context, + ), ) + assert isinstance(result.messages[4], TextMessage) assert result.messages[4].content == "Hello" assert result.messages[4].source == "agent2" + assert isinstance(result.messages[5], TextMessage) assert result.messages[5].content == "TERMINATE" assert result.messages[5].source == "agent2" @@ -1020,17 +1626,26 @@ async def test_swarm_with_handoff_termination(runtime: AgentRuntime | None) -> N # Start result = await team.run(task="task") assert len(result.messages) == 2 + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result.messages[1], HandoffMessage) assert result.messages[0].content == "task" assert result.messages[1].content == "Transferred to third_agent." # Resume existing. result = await team.run() assert len(result.messages) == 3 + assert isinstance(result.messages[0], HandoffMessage) + assert isinstance(result.messages[1], HandoffMessage) + assert isinstance(result.messages[2], HandoffMessage) assert result.messages[0].content == "Transferred to first_agent." assert result.messages[1].content == "Transferred to second_agent." assert result.messages[2].content == "Transferred to third_agent." # Resume new task. result = await team.run(task="new task") assert len(result.messages) == 4 + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result.messages[1], HandoffMessage) + assert isinstance(result.messages[2], HandoffMessage) + assert isinstance(result.messages[3], HandoffMessage) assert result.messages[0].content == "new task" assert result.messages[1].content == "Transferred to first_agent." assert result.messages[2].content == "Transferred to second_agent." @@ -1043,6 +1658,9 @@ async def test_swarm_with_handoff_termination(runtime: AgentRuntime | None) -> N # Start result = await team.run(task="task") assert len(result.messages) == 3 + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result.messages[1], HandoffMessage) + assert isinstance(result.messages[2], HandoffMessage) assert result.messages[0].content == "task" assert result.messages[1].content == "Transferred to third_agent." assert result.messages[2].content == "Transferred to non_existing_agent." @@ -1055,6 +1673,10 @@ async def test_swarm_with_handoff_termination(runtime: AgentRuntime | None) -> N # Resume with a HandoffMessage result = await team.run(task=HandoffMessage(content="Handoff to first_agent.", target="first_agent", source="user")) assert len(result.messages) == 4 + assert isinstance(result.messages[0], HandoffMessage) + assert isinstance(result.messages[1], HandoffMessage) + assert isinstance(result.messages[2], HandoffMessage) + assert isinstance(result.messages[3], HandoffMessage) assert result.messages[0].content == "Handoff to first_agent." assert result.messages[1].content == "Transferred to second_agent." assert result.messages[2].content == "Transferred to third_agent." @@ -1070,7 +1692,7 @@ async def test_round_robin_group_chat_with_message_list(runtime: AgentRuntime | team = RoundRobinGroupChat([agent1, agent2], termination_condition=termination, runtime=runtime) # Create a list of messages - messages: List[ChatMessage] = [ + messages: List[BaseChatMessage] = [ TextMessage(content="Message 1", source="user"), TextMessage(content="Message 2", source="user"), TextMessage(content="Message 3", source="user"), @@ -1081,6 +1703,10 @@ async def test_round_robin_group_chat_with_message_list(runtime: AgentRuntime | # Verify the messages were processed in order assert len(result.messages) == 4 # Initial messages + echo until termination + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result.messages[1], TextMessage) + assert isinstance(result.messages[2], TextMessage) + assert isinstance(result.messages[3], TextMessage) assert result.messages[0].content == "Message 1" # First message assert result.messages[1].content == "Message 2" # Second message assert result.messages[2].content == "Message 3" # Third message @@ -1089,16 +1715,16 @@ async def test_round_robin_group_chat_with_message_list(runtime: AgentRuntime | # Test with streaming await team.reset() - index = 0 + result_index = 0 # Include the 3 task messages in result since output_task_messages=True by default async for message in team.run_stream(task=messages): if isinstance(message, TaskResult): - assert message == result + assert compare_task_results(message, result) else: - assert message == result.messages[index] - index += 1 + assert compare_messages(message, result.messages[result_index]) + result_index += 1 # Test with invalid message list - with pytest.raises(ValueError, match="All messages in task list must be valid ChatMessage types"): + with pytest.raises(ValueError, match="All messages in task list must be valid BaseChatMessage types"): await team.run(task=["not a message"]) # type: ignore[list-item, arg-type] # intentionally testing invalid input # Test with empty message list @@ -1111,12 +1737,14 @@ async def test_declarative_groupchats_with_config(runtime: AgentRuntime | None) # Create basic agents and components for testing agent1 = AssistantAgent( "agent_1", - model_client=OpenAIChatCompletionClient(model="gpt-4o-2024-05-13", api_key=""), + model_client=OpenAIChatCompletionClient(model="gpt-4.1-nano-2025-04-14", api_key=""), handoffs=["agent_2"], ) - agent2 = AssistantAgent("agent_2", model_client=OpenAIChatCompletionClient(model="gpt-4o-2024-05-13", api_key="")) + agent2 = AssistantAgent( + "agent_2", model_client=OpenAIChatCompletionClient(model="gpt-4.1-nano-2025-04-14", api_key="") + ) termination = MaxMessageTermination(4) - model_client = OpenAIChatCompletionClient(model="gpt-4o-2024-05-13", api_key="") + model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano-2025-04-14", api_key="") # Test round robin - verify config is preserved round_robin = RoundRobinGroupChat(participants=[agent1, agent2], termination_condition=termination, max_turns=5) @@ -1170,3 +1798,149 @@ async def test_declarative_groupchats_with_config(runtime: AgentRuntime | None) assert selector.dump_component().provider == "autogen_agentchat.teams.SelectorGroupChat" assert swarm.dump_component().provider == "autogen_agentchat.teams.Swarm" assert magentic.dump_component().provider == "autogen_agentchat.teams.MagenticOneGroupChat" + + +class _StructuredContent(BaseModel): + message: str + + +class _StructuredAgent(BaseChatAgent): + def __init__(self, name: str, description: str) -> None: + super().__init__(name, description) + self._message = _StructuredContent(message="Structured hello") + + @property + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: + return (StructuredMessage[_StructuredContent],) + + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: + return Response( + chat_message=StructuredMessage[_StructuredContent]( + source=self.name, + content=self._message, + format_string="Structured says: {message}", + ) + ) + + async def on_reset(self, cancellation_token: CancellationToken) -> None: + pass + + +@pytest.mark.asyncio +async def test_message_type_auto_registration(runtime: AgentRuntime | None) -> None: + agent1 = _StructuredAgent("structured", description="emits structured messages") + agent2 = _EchoAgent("echo", description="echoes input") + + team = RoundRobinGroupChat(participants=[agent1, agent2], max_turns=2, runtime=runtime) + + result = await team.run(task="Say something structured") + + assert len(result.messages) == 3 + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result.messages[1], StructuredMessage) + assert isinstance(result.messages[2], TextMessage) + assert result.messages[1].to_text() == "Structured says: Structured hello" + + +@pytest.mark.asyncio +async def test_structured_message_state_roundtrip(runtime: AgentRuntime | None) -> None: + agent1 = _StructuredAgent("structured", description="sends structured") + agent2 = _EchoAgent("echo", description="echoes") + + team1 = RoundRobinGroupChat( + participants=[agent1, agent2], + termination_condition=MaxMessageTermination(2), + runtime=runtime, + ) + + await team1.run(task="Say something structured") + state1 = await team1.save_state() + + # Recreate team without needing custom_message_types + agent3 = _StructuredAgent("structured", description="sends structured") + agent4 = _EchoAgent("echo", description="echoes") + team2 = RoundRobinGroupChat( + participants=[agent3, agent4], + termination_condition=MaxMessageTermination(2), + runtime=runtime, + ) + + await team2.load_state(state1) + state2 = await team2.save_state() + + # Assert full state equality + assert state1 == state2 + + # Assert message thread content match + manager1 = await team1._runtime.try_get_underlying_agent_instance( # pyright: ignore + AgentId(f"{team1._group_chat_manager_name}_{team1._team_id}", team1._team_id), # pyright: ignore + RoundRobinGroupChatManager, + ) + manager2 = await team2._runtime.try_get_underlying_agent_instance( # pyright: ignore + AgentId(f"{team2._group_chat_manager_name}_{team2._team_id}", team2._team_id), # pyright: ignore + RoundRobinGroupChatManager, + ) + + assert manager1._message_thread == manager2._message_thread # pyright: ignore + + +@pytest.mark.asyncio +async def test_selector_group_chat_streaming(runtime: AgentRuntime | None) -> None: + model_client = ReplayChatCompletionClient( + ["the agent should be agent2"], + ) + agent2 = _StopAgent("agent2", description="stop agent 2", stop_at=0) + agent3 = _EchoAgent("agent3", description="echo agent 3") + termination = StopMessageTermination() + team = SelectorGroupChat( + participants=[agent2, agent3], + model_client=model_client, + termination_condition=termination, + runtime=runtime, + emit_team_events=True, + model_client_streaming=True, + ) + result = await team.run( + task="Write a program that prints 'Hello, world!'", + ) + + assert len(result.messages) == 4 + assert isinstance(result.messages[0], TextMessage) + assert isinstance(result.messages[1], SelectorEvent) + assert isinstance(result.messages[2], SelectSpeakerEvent) + assert isinstance(result.messages[3], StopMessage) + + assert result.messages[0].content == "Write a program that prints 'Hello, world!'" + assert result.messages[1].content == "the agent should be agent2" + assert result.messages[2].content == ["agent2"] + assert result.messages[3].source == "agent2" + assert result.stop_reason is not None and result.stop_reason == "Stop message received" + + # Test streaming + await team.reset() + model_client.reset() + result_index = 0 # Include task message in result since output_task_messages=True by default + streamed_chunks: List[str] = [] + final_result: TaskResult | None = None + async for message in team.run_stream( + task="Write a program that prints 'Hello, world!'", + ): + if isinstance(message, TaskResult): + final_result = message + assert compare_task_results(message, result) + elif isinstance(message, ModelClientStreamingChunkEvent): + streamed_chunks.append(message.content) + else: + if streamed_chunks: + assert isinstance(message, SelectorEvent) + assert message.content == "".join(streamed_chunks) + streamed_chunks = [] + assert compare_messages(message, result.messages[result_index]) + result_index += 1 + + # Verify we got the expected messages without relying on fragile ordering + assert final_result is not None + assert len(streamed_chunks) == 0 # All chunks should have been processed + + # Content-based verification instead of index-based + # Note: The streaming test verifies the streaming behavior, not the final result content diff --git a/python/packages/autogen-agentchat/tests/test_group_chat_endpoint.py b/python/packages/autogen-agentchat/tests/test_group_chat_endpoint.py index 390a45e031f4..142df272950f 100644 --- a/python/packages/autogen-agentchat/tests/test_group_chat_endpoint.py +++ b/python/packages/autogen-agentchat/tests/test_group_chat_endpoint.py @@ -4,10 +4,7 @@ import pytest from autogen_agentchat.agents import AssistantAgent from autogen_agentchat.base import TaskResult -from autogen_agentchat.messages import ( - AgentEvent, - ChatMessage, -) +from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage from autogen_agentchat.teams import SelectorGroupChat from autogen_agentchat.ui import Console from autogen_core.models import ChatCompletionClient @@ -36,7 +33,7 @@ async def _test_selector_group_chat(model_client: ChatCompletionClient) -> None: async def _test_selector_group_chat_with_candidate_func(model_client: ChatCompletionClient) -> None: filtered_participants = ["developer", "tester"] - def dummy_candidate_func(thread: Sequence[AgentEvent | ChatMessage]) -> List[str]: + def dummy_candidate_func(thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str]: # Dummy candidate function that will return # only return developer and reviewer return filtered_participants @@ -101,7 +98,7 @@ async def test_selector_group_chat_openai() -> None: pytest.skip("OPENAI_API_KEY not set in environment variables.") model_client = OpenAIChatCompletionClient( - model="gpt-4o-mini", + model="gpt-4.1-nano", api_key=api_key, ) await _test_selector_group_chat(model_client) diff --git a/python/packages/autogen-agentchat/tests/test_group_chat_graph.py b/python/packages/autogen-agentchat/tests/test_group_chat_graph.py new file mode 100644 index 000000000000..98cc8ca66b34 --- /dev/null +++ b/python/packages/autogen-agentchat/tests/test_group_chat_graph.py @@ -0,0 +1,1723 @@ +import asyncio +import re +from typing import AsyncGenerator, List, Sequence +from unittest.mock import patch + +import pytest +import pytest_asyncio +from autogen_agentchat.agents import ( + AssistantAgent, + BaseChatAgent, + MessageFilterAgent, + MessageFilterConfig, + PerSourceFilter, +) +from autogen_agentchat.base import Response, TaskResult +from autogen_agentchat.conditions import MaxMessageTermination, SourceMatchTermination +from autogen_agentchat.messages import BaseChatMessage, ChatMessage, MessageFactory, StopMessage, TextMessage +from autogen_agentchat.teams import ( + DiGraphBuilder, + GraphFlow, +) +from autogen_agentchat.teams._group_chat._events import ( # type: ignore[attr-defined] + BaseAgentEvent, + GroupChatTermination, +) +from autogen_agentchat.teams._group_chat._graph._digraph_group_chat import ( + DiGraph, + DiGraphEdge, + DiGraphNode, + GraphFlowManager, +) +from autogen_core import AgentRuntime, CancellationToken, Component, SingleThreadedAgentRuntime +from autogen_ext.models.replay import ReplayChatCompletionClient +from pydantic import BaseModel +from utils import compare_message_lists, compare_task_results + + +def test_create_digraph() -> None: + """Test creating a simple directed graph.""" + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + + assert "A" in graph.nodes + assert "B" in graph.nodes + assert "C" in graph.nodes + assert len(graph.nodes["A"].edges) == 1 + assert len(graph.nodes["B"].edges) == 1 + assert len(graph.nodes["C"].edges) == 0 + + +def test_get_parents() -> None: + """Test computing parent relationships.""" + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + + parents = graph.get_parents() + assert parents["A"] == [] + assert parents["B"] == ["A"] + assert parents["C"] == ["B"] + + +def test_get_start_nodes() -> None: + """Test retrieving start nodes (nodes with no incoming edges).""" + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + + start_nodes = graph.get_start_nodes() + assert start_nodes == set(["A"]) + + +def test_get_leaf_nodes() -> None: + """Test retrieving leaf nodes (nodes with no outgoing edges).""" + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + + leaf_nodes = graph.get_leaf_nodes() + assert leaf_nodes == set(["C"]) + + +def test_serialization() -> None: + """Test serializing and deserializing the graph.""" + # Use a string condition instead of a lambda + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B", condition="trigger1")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + + serialized = graph.model_dump_json() + deserialized_graph = DiGraph.model_validate_json(serialized) + + assert deserialized_graph.nodes["A"].edges[0].target == "B" + assert deserialized_graph.nodes["A"].edges[0].condition == "trigger1" + assert deserialized_graph.nodes["B"].edges[0].target == "C" + + # Test the original condition works + test_msg = TextMessage(content="this has trigger1 in it", source="test") + # Manually check if the string is in the message text + assert "trigger1" in test_msg.to_model_text() + + +def test_invalid_graph_no_start_node() -> None: + """Test validation failure when there is no start node.""" + graph = DiGraph( + nodes={ + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[DiGraphEdge(target="B")]), # Forms a cycle + } + ) + + start_nodes = graph.get_start_nodes() + assert len(start_nodes) == 0 # Now it correctly fails when no start nodes exist + + +def test_invalid_graph_no_leaf_node() -> None: + """Test validation failure when there is no leaf node.""" + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[DiGraphEdge(target="A")]), # Circular reference + } + ) + + leaf_nodes = graph.get_leaf_nodes() + assert len(leaf_nodes) == 0 # No true endpoint because of cycle + + +def test_condition_edge_execution() -> None: + """Test conditional edge execution support.""" + # Use string condition + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B", condition="TRIGGER")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + + # Check the condition manually + test_message = TextMessage(content="This has TRIGGER in it", source="test") + non_match_message = TextMessage(content="This doesn't match", source="test") + + # Check if the string condition is in each message text + assert "TRIGGER" in test_message.to_model_text() + assert "TRIGGER" not in non_match_message.to_model_text() + + # Check the condition itself + assert graph.nodes["A"].edges[0].condition == "TRIGGER" + assert graph.nodes["B"].edges[0].condition is None + + +def test_graph_with_multiple_paths() -> None: + """Test a graph with multiple execution paths.""" + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B"), DiGraphEdge(target="C")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="D")]), + "C": DiGraphNode(name="C", edges=[DiGraphEdge(target="D")]), + "D": DiGraphNode(name="D", edges=[]), + } + ) + + parents = graph.get_parents() + assert parents["B"] == ["A"] + assert parents["C"] == ["A"] + assert parents["D"] == ["B", "C"] + + start_nodes = graph.get_start_nodes() + assert start_nodes == set(["A"]) + + leaf_nodes = graph.get_leaf_nodes() + assert leaf_nodes == set(["D"]) + + +def test_cycle_detection_no_cycle() -> None: + """Test that a valid acyclic graph returns False for cycle check.""" + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + assert not graph.has_cycles_with_exit() + + +def test_cycle_detection_with_exit_condition() -> None: + """Test a graph with cycle and conditional exit passes validation.""" + # Use a string condition + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[DiGraphEdge(target="A", condition="exit")]), # Cycle with condition + } + ) + assert graph.has_cycles_with_exit() + + # Use a lambda condition + graph_with_lambda = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode( + name="C", edges=[DiGraphEdge(target="A", condition=lambda msg: "test" in msg.to_model_text())] + ), # Cycle with lambda + } + ) + assert graph_with_lambda.has_cycles_with_exit() + + +def test_cycle_detection_without_exit_condition() -> None: + """Test that cycle without exit condition raises an error.""" + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[DiGraphEdge(target="A")]), # Cycle without condition + "D": DiGraphNode(name="D", edges=[DiGraphEdge(target="E")]), + "E": DiGraphNode(name="E", edges=[]), + } + ) + with pytest.raises(ValueError, match="Cycle detected without exit condition: A -> B -> C -> A"): + graph.has_cycles_with_exit() + + +def test_different_activation_groups_detection() -> None: + """Test different activation groups.""" + graph = DiGraph( + nodes={ + "A": DiGraphNode( + name="A", + edges=[ + DiGraphEdge(target="B"), + DiGraphEdge(target="C"), + ], + ), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="D", activation_condition="all")]), + "C": DiGraphNode(name="C", edges=[DiGraphEdge(target="D", activation_condition="any")]), + "D": DiGraphNode(name="D", edges=[]), + } + ) + with pytest.raises( + ValueError, + match=re.escape( + "Conflicting activation conditions for target 'D' group 'D': " + "'all' (from node 'B') and 'any' (from node 'C')" + ), + ): + graph.graph_validate() + + +def test_validate_graph_success() -> None: + """Test successful validation of a valid graph.""" + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[]), + } + ) + # No error should be raised + graph.graph_validate() + assert not graph.get_has_cycles() + + # Use a lambda condition + graph_with_lambda = DiGraph( + nodes={ + "A": DiGraphNode( + name="A", edges=[DiGraphEdge(target="B", condition=lambda msg: "test" in msg.to_model_text())] + ), + "B": DiGraphNode(name="B", edges=[]), + } + ) + # No error should be raised + graph_with_lambda.graph_validate() + assert not graph_with_lambda.get_has_cycles() + + +def test_validate_graph_missing_start_node() -> None: + """Test validation failure when no start node exists.""" + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="A")]), # Cycle + } + ) + with pytest.raises(ValueError, match="Graph must have at least one start node"): + graph.graph_validate() + + +def test_validate_graph_missing_leaf_node() -> None: + """Test validation failure when no leaf node exists.""" + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[DiGraphEdge(target="B")]), # Cycle + } + ) + with pytest.raises(ValueError, match="Graph must have at least one leaf node"): + graph.graph_validate() + + +def test_validate_graph_mixed_conditions() -> None: + """Test validation failure when node has mixed conditional and unconditional edges.""" + # Use string for condition + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B", condition="cond"), DiGraphEdge(target="C")]), + "B": DiGraphNode(name="B", edges=[]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + with pytest.raises(ValueError, match="Node 'A' has a mix of conditional and unconditional edges"): + graph.graph_validate() + + # Use lambda for condition + graph_with_lambda = DiGraph( + nodes={ + "A": DiGraphNode( + name="A", + edges=[ + DiGraphEdge(target="B", condition=lambda msg: "test" in msg.to_model_text()), + DiGraphEdge(target="C"), + ], + ), + "B": DiGraphNode(name="B", edges=[]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + with pytest.raises(ValueError, match="Node 'A' has a mix of conditional and unconditional edges"): + graph_with_lambda.graph_validate() + + +@pytest.mark.asyncio +async def test_invalid_digraph_manager_cycle_without_termination() -> None: + """Test GraphManager raises error for cyclic graph without termination condition.""" + # Create a cyclic graph A → B → A + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="A")]), + } + ) + + output_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination] = asyncio.Queue() + + with patch( + "autogen_agentchat.teams._group_chat._base_group_chat_manager.BaseGroupChatManager.__init__", + return_value=None, + ): + manager = GraphFlowManager.__new__(GraphFlowManager) + + with pytest.raises(ValueError, match="Graph must have at least one start node"): + manager.__init__( # type: ignore[misc] + name="test_manager", + group_topic_type="topic", + output_topic_type="topic", + participant_topic_types=["topic1", "topic2"], + participant_names=["A", "B"], + participant_descriptions=["Agent A", "Agent B"], + output_message_queue=output_queue, + termination_condition=None, + max_turns=None, + message_factory=MessageFactory(), + graph=graph, + ) + + +class _EchoAgent(BaseChatAgent): + def __init__(self, name: str, description: str) -> None: + super().__init__(name, description) + self._last_message: str | None = None + self._total_messages = 0 + + @property + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: + return (TextMessage,) + + @property + def total_messages(self) -> int: + return self._total_messages + + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: + if len(messages) > 0: + assert isinstance(messages[0], TextMessage) or isinstance(messages[0], StopMessage) + self._last_message = messages[0].content + self._total_messages += 1 + return Response(chat_message=TextMessage(content=messages[0].content, source=self.name)) + else: + assert self._last_message is not None + self._total_messages += 1 + return Response(chat_message=TextMessage(content=self._last_message, source=self.name)) + + async def on_reset(self, cancellation_token: CancellationToken) -> None: + self._last_message = None + + +@pytest_asyncio.fixture(params=["single_threaded", "embedded"]) # type: ignore +async def runtime(request: pytest.FixtureRequest) -> AsyncGenerator[AgentRuntime | None, None]: + if request.param == "single_threaded": + runtime = SingleThreadedAgentRuntime() + runtime.start() + yield runtime + await runtime.stop() + elif request.param == "embedded": + yield None + + +TaskType = str | List[ChatMessage] | ChatMessage + + +@pytest.mark.asyncio +async def test_digraph_group_chat_sequential_execution(runtime: AgentRuntime | None) -> None: + # Create agents A → B → C + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + + # Define graph A → B → C + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + + # Create team using Graph + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + + # Run the chat + result: TaskResult = await team.run(task="Hello from User") + + assert len(result.messages) == 4 + assert isinstance(result.messages[0], TextMessage) + assert result.messages[0].source == "user" + assert result.messages[1].source == "A" + assert result.messages[2].source == "B" + assert result.messages[3].source == "C" + assert all(isinstance(m, TextMessage) for m in result.messages) + assert result.stop_reason is not None + + +@pytest.mark.asyncio +async def test_digraph_group_chat_parallel_fanout(runtime: AgentRuntime | None) -> None: + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B"), DiGraphEdge(target="C")]), + "B": DiGraphNode(name="B", edges=[]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + + result: TaskResult = await team.run(task="Start") + assert len(result.messages) == 4 + assert result.messages[0].source == "user" + assert result.messages[1].source == "A" + assert set(m.source for m in result.messages[2:]) == {"B", "C"} + assert result.stop_reason is not None + + +@pytest.mark.asyncio +async def test_digraph_group_chat_parallel_join_all(runtime: AgentRuntime | None) -> None: + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="C")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[], activation="all"), + } + ) + + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + + result: TaskResult = await team.run(task="Go") + assert len(result.messages) == 4 + assert result.messages[0].source == "user" + assert set([result.messages[1].source, result.messages[2].source]) == {"A", "B"} + assert result.messages[3].source == "C" + assert result.stop_reason is not None + + +@pytest.mark.asyncio +async def test_digraph_group_chat_parallel_join_any(runtime: AgentRuntime | None) -> None: + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="C")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[], activation="any"), + } + ) + + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + + result: TaskResult = await team.run(task="Start") + + assert len(result.messages) == 4 + assert result.messages[0].source == "user" + sources = [m.source for m in result.messages[1:]] + + # C must be last + assert sources[-1] == "C" + + # A and B must both execute + assert {"A", "B"}.issubset(set(sources)) + + # One of A or B must execute before C + index_a = sources.index("A") + index_b = sources.index("B") + index_c = sources.index("C") + assert index_c > min(index_a, index_b) + assert result.stop_reason is not None + + +@pytest.mark.asyncio +async def test_digraph_group_chat_multiple_start_nodes(runtime: AgentRuntime | None) -> None: + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[]), + "B": DiGraphNode(name="B", edges=[]), + } + ) + + team = GraphFlow( + participants=[agent_a, agent_b], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + + result: TaskResult = await team.run(task="Start") + assert len(result.messages) == 3 + assert result.messages[0].source == "user" + assert set(m.source for m in result.messages[1:]) == {"A", "B"} + assert result.stop_reason is not None + + +@pytest.mark.asyncio +async def test_digraph_group_chat_disconnected_graph(runtime: AgentRuntime | None) -> None: + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + agent_d = _EchoAgent("D", description="Echo agent D") + + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[]), + "C": DiGraphNode(name="C", edges=[DiGraphEdge(target="D")]), + "D": DiGraphNode(name="D", edges=[]), + } + ) + + team = GraphFlow( + participants=[agent_a, agent_b, agent_c, agent_d], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(10), + ) + + result: TaskResult = await team.run(task="Go") + assert len(result.messages) == 5 + assert result.messages[0].source == "user" + assert {"A", "C"} == set([result.messages[1].source, result.messages[2].source]) + assert {"B", "D"} == set([result.messages[3].source, result.messages[4].source]) + assert result.stop_reason is not None + + +@pytest.mark.asyncio +async def test_digraph_group_chat_conditional_branch(runtime: AgentRuntime | None) -> None: + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + + # Use string conditions + graph = DiGraph( + nodes={ + "A": DiGraphNode( + name="A", edges=[DiGraphEdge(target="B", condition="yes"), DiGraphEdge(target="C", condition="no")] + ), + "B": DiGraphNode(name="B", edges=[], activation="any"), + "C": DiGraphNode(name="C", edges=[], activation="any"), + } + ) + + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + + result = await team.run(task="Trigger yes") + assert result.messages[2].source == "B" + + # Use lambda conditions + graph_with_lambda = DiGraph( + nodes={ + "A": DiGraphNode( + name="A", + edges=[ + DiGraphEdge(target="B", condition=lambda msg: "yes" in msg.to_model_text()), + DiGraphEdge(target="C", condition=lambda msg: "no" in msg.to_model_text()), + ], + ), + "B": DiGraphNode(name="B", edges=[], activation="any"), + "C": DiGraphNode(name="C", edges=[], activation="any"), + } + ) + team_with_lambda = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph_with_lambda, + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + result_with_lambda = await team_with_lambda.run(task="Trigger no") + assert result_with_lambda.messages[2].source == "C" + + +@pytest.mark.asyncio +async def test_digraph_group_chat_loop_with_exit_condition(runtime: AgentRuntime | None) -> None: + # Agents A and C: Echo Agents + agent_a = _EchoAgent("A", description="Echo agent A") + agent_c = _EchoAgent("C", description="Echo agent C") + + # Replay model client for agent B + model_client = ReplayChatCompletionClient( + chat_completions=[ + "loop", # First time B will ask to loop + "loop", # Second time B will ask to loop + "exit", # Third time B will say exit + ] + ) + # Agent B: Assistant Agent using Replay Client + agent_b = AssistantAgent("B", description="Decision agent B", model_client=model_client) + + # DiGraph: A → B → C (conditional back to A or terminate) + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode( + name="B", edges=[DiGraphEdge(target="C", condition="exit"), DiGraphEdge(target="A", condition="loop")] + ), + "C": DiGraphNode(name="C", edges=[]), + }, + default_start_node="A", + ) + + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(20), + ) + + # Run + result = await team.run(task="Start") + + # Assert message order + expected_sources = [ + "user", + "A", + "B", # 1st loop + "A", + "B", # 2nd loop + "A", + "B", + "C", + ] + + actual_sources = [m.source for m in result.messages] + + assert actual_sources == expected_sources + assert result.stop_reason is not None + assert result.messages[-1].source == "C" + assert any(m.content == "exit" for m in result.messages) # type: ignore[attr-defined,union-attr] + + +@pytest.mark.asyncio +async def test_digraph_group_chat_loop_with_self_cycle(runtime: AgentRuntime | None) -> None: + # Agents A and C: Echo Agents + agent_a = _EchoAgent("A", description="Echo agent A") + agent_c = _EchoAgent("C", description="Echo agent C") + + # Replay model client for agent B + model_client = ReplayChatCompletionClient( + chat_completions=[ + "loop", # First time B will ask to loop + "loop", # Second time B will ask to loop + "exit", # Third time B will say exit + ] + ) + # Agent B: Assistant Agent using Replay Client + agent_b = AssistantAgent("B", description="Decision agent B", model_client=model_client) + + # DiGraph: A → B(self loop) → C (conditional back to A or terminate) + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode( + name="B", + edges=[ + DiGraphEdge(target="C", condition="exit"), + DiGraphEdge(target="B", condition="loop", activation_group="B_loop"), + ], + ), + "C": DiGraphNode(name="C", edges=[]), + }, + default_start_node="A", + ) + + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(20), + ) + + # Run + result = await team.run(task="Start") + + # Assert message order + expected_sources = [ + "user", + "A", + "B", # 1st loop + "B", # 2nd loop + "B", + "C", + ] + + actual_sources = [m.source for m in result.messages] + + assert actual_sources == expected_sources + assert result.stop_reason is not None + assert result.messages[-1].source == "C" + assert any(m.content == "exit" for m in result.messages) # type: ignore[attr-defined,union-attr] + + +@pytest.mark.asyncio +async def test_digraph_group_chat_loop_with_two_cycles(runtime: AgentRuntime | None) -> None: + # Agents A and C: Echo Agents + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + agent_e = _EchoAgent("E", description="Echo agent E") + + # Replay model client for agent B + model_client = ReplayChatCompletionClient( + chat_completions=[ + "to_x", # First time O will branch to B + "to_o", # X will go back to O + "to_y", # Second time O will branch to C + "to_o", # Y will go back to O + "exit", # Third time O will say exit + ] + ) + # Agent o, b, c: Assistant Agent using Replay Client + agent_o = AssistantAgent("O", description="Decision agent o", model_client=model_client) + agent_x = AssistantAgent("X", description="Decision agent x", model_client=model_client) + agent_y = AssistantAgent("Y", description="Decision agent y", model_client=model_client) + + # DiGraph: + # + # A + # / \ + # B C + # \ | + # X = O = Y (bidirectional) + # | + # E(exit) + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B"), DiGraphEdge(target="C")]), + "B": DiGraphNode( + name="B", edges=[DiGraphEdge(target="O")] + ), # default activation group name is same as target node name "O" + "C": DiGraphNode( + name="C", edges=[DiGraphEdge(target="O")] + ), # default activation group name is same as target node name "O" + "O": DiGraphNode( + name="O", + edges=[ + DiGraphEdge(target="X", condition="to_x"), + DiGraphEdge(target="Y", condition="to_y"), + DiGraphEdge(target="E", condition="exit"), + ], + ), + "X": DiGraphNode(name="X", edges=[DiGraphEdge(target="O", condition="to_o", activation_group="x_o_loop")]), + "Y": DiGraphNode(name="Y", edges=[DiGraphEdge(target="O", condition="to_o", activation_group="y_o_loop")]), + "E": DiGraphNode(name="E", edges=[]), + }, + default_start_node="A", + ) + + team = GraphFlow( + participants=[agent_a, agent_o, agent_b, agent_c, agent_x, agent_y, agent_e], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(20), + ) + + # Run + result = await team.run(task="Start") + + # Assert message order + expected_sources = [ + "user", + "A", + "B", + "C", + "O", + "X", # O -> X + "O", # X -> O + "Y", # O -> Y + "O", # Y -> O + "E", # O -> E + ] + + actual_sources = [m.source for m in result.messages] + + assert actual_sources == expected_sources + assert result.stop_reason is not None + assert result.messages[-1].source == "E" + assert any(m.content == "exit" for m in result.messages) # type: ignore[attr-defined,union-attr] + + +@pytest.mark.asyncio +async def test_digraph_group_chat_parallel_join_any_1(runtime: AgentRuntime | None) -> None: + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + agent_d = _EchoAgent("D", description="Echo agent D") + + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B"), DiGraphEdge(target="C")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="D", activation_group="any")]), + "C": DiGraphNode(name="C", edges=[DiGraphEdge(target="D", activation_group="any")]), + "D": DiGraphNode(name="D", edges=[]), + } + ) + + team = GraphFlow( + participants=[agent_a, agent_b, agent_c, agent_d], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(10), + ) + + result = await team.run(task="Run parallel join") + sequence = [msg.source for msg in result.messages if isinstance(msg, TextMessage)] + assert sequence[0] == "user" + # B and C should both run + assert "B" in sequence + assert "C" in sequence + # D should trigger twice → once after B and once after C (order depends on runtime) + d_indices = [i for i, s in enumerate(sequence) if s == "D"] + assert len(d_indices) == 1 + # Each D trigger must be after corresponding B or C + b_index = sequence.index("B") + c_index = sequence.index("C") + assert any(d > b_index for d in d_indices) + assert any(d > c_index for d in d_indices) + assert result.stop_reason is not None + + +@pytest.mark.asyncio +async def test_digraph_group_chat_chained_parallel_join_any(runtime: AgentRuntime | None) -> None: + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + agent_d = _EchoAgent("D", description="Echo agent D") + agent_e = _EchoAgent("E", description="Echo agent E") + + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B"), DiGraphEdge(target="C")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="D")]), + "C": DiGraphNode(name="C", edges=[DiGraphEdge(target="D")]), + "D": DiGraphNode(name="D", edges=[DiGraphEdge(target="E")], activation="any"), + "E": DiGraphNode(name="E", edges=[], activation="any"), + } + ) + + team = GraphFlow( + participants=[agent_a, agent_b, agent_c, agent_d, agent_e], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(20), + ) + + result = await team.run(task="Run chained parallel join-any") + + sequence = [msg.source for msg in result.messages if isinstance(msg, TextMessage)] + + # D should trigger twice + d_indices = [i for i, s in enumerate(sequence) if s == "D"] + assert len(d_indices) == 1 + # Each D trigger must be after corresponding B or C + b_index = sequence.index("B") + c_index = sequence.index("C") + assert any(d > b_index for d in d_indices) + assert any(d > c_index for d in d_indices) + + # E should also trigger twice → once after each D + e_indices = [i for i, s in enumerate(sequence) if s == "E"] + assert len(e_indices) == 1 + assert e_indices[0] > d_indices[0] + assert result.stop_reason is not None + + +@pytest.mark.asyncio +async def test_digraph_group_chat_multiple_conditional(runtime: AgentRuntime | None) -> None: + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + agent_d = _EchoAgent("D", description="Echo agent D") + + # Use string conditions + graph = DiGraph( + nodes={ + "A": DiGraphNode( + name="A", + edges=[ + DiGraphEdge(target="B", condition="apple"), + DiGraphEdge(target="C", condition="banana"), + DiGraphEdge(target="D", condition="cherry"), + ], + ), + "B": DiGraphNode(name="B", edges=[]), + "C": DiGraphNode(name="C", edges=[]), + "D": DiGraphNode(name="D", edges=[]), + } + ) + + team = GraphFlow( + participants=[agent_a, agent_b, agent_c, agent_d], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + + # Test banana branch + result = await team.run(task="banana") + assert result.messages[2].source == "C" + + # Use lambda conditions + graph_with_lambda = DiGraph( + nodes={ + "A": DiGraphNode( + name="A", + edges=[ + DiGraphEdge(target="B", condition=lambda msg: "apple" in msg.to_model_text()), + DiGraphEdge(target="C", condition=lambda msg: "banana" in msg.to_model_text()), + DiGraphEdge(target="D", condition=lambda msg: "cherry" in msg.to_model_text()), + ], + ), + "B": DiGraphNode(name="B", edges=[]), + "C": DiGraphNode(name="C", edges=[]), + "D": DiGraphNode(name="D", edges=[]), + } + ) + team_with_lambda = GraphFlow( + participants=[agent_a, agent_b, agent_c, agent_d], + graph=graph_with_lambda, + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + result_with_lambda = await team_with_lambda.run(task="cherry") + assert result_with_lambda.messages[2].source == "D" + + +class _TestMessageFilterAgentConfig(BaseModel): + name: str + description: str = "Echo test agent" + + +class _TestMessageFilterAgent(BaseChatAgent, Component[_TestMessageFilterAgentConfig]): + component_config_schema = _TestMessageFilterAgentConfig + component_provider_override = "test_group_chat_graph._TestMessageFilterAgent" + + def __init__(self, name: str, description: str = "Echo test agent") -> None: + super().__init__(name=name, description=description) + self.received_messages: list[BaseChatMessage] = [] + + @property + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: + return (TextMessage,) + + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: + self.received_messages.extend(messages) + return Response(chat_message=TextMessage(content="ACK", source=self.name)) + + async def on_reset(self, cancellation_token: CancellationToken) -> None: + self.received_messages.clear() + + def _to_config(self) -> _TestMessageFilterAgentConfig: + return _TestMessageFilterAgentConfig(name=self.name, description=self.description) + + @classmethod + def _from_config(cls, config: _TestMessageFilterAgentConfig) -> "_TestMessageFilterAgent": + return cls(name=config.name, description=config.description) + + +@pytest.mark.asyncio +async def test_message_filter_agent_empty_filter_blocks_all() -> None: + inner_agent = _TestMessageFilterAgent("inner") + wrapper = MessageFilterAgent( + name="wrapper", + wrapped_agent=inner_agent, + filter=MessageFilterConfig(per_source=[]), + ) + messages = [ + TextMessage(source="user", content="Hello"), + TextMessage(source="system", content="System msg"), + ] + await wrapper.on_messages(messages, CancellationToken()) + assert len(inner_agent.received_messages) == 0 + + +@pytest.mark.asyncio +async def test_message_filter_agent_with_position_none_gets_all() -> None: + inner_agent = _TestMessageFilterAgent("inner") + wrapper = MessageFilterAgent( + name="wrapper", + wrapped_agent=inner_agent, + filter=MessageFilterConfig(per_source=[PerSourceFilter(source="user", position=None, count=None)]), + ) + messages = [ + TextMessage(source="user", content="A"), + TextMessage(source="user", content="B"), + TextMessage(source="system", content="Ignore this"), + ] + await wrapper.on_messages(messages, CancellationToken()) + assert len(inner_agent.received_messages) == 2 + assert {m.content for m in inner_agent.received_messages} == {"A", "B"} # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_digraph_group_chat() -> None: + inner_agent = _TestMessageFilterAgent("agent") + wrapper = MessageFilterAgent( + name="agent", + wrapped_agent=inner_agent, + filter=MessageFilterConfig( + per_source=[ + PerSourceFilter(source="user", position="last", count=2), + PerSourceFilter(source="system", position="first", count=1), + ] + ), + ) + config = wrapper.dump_component() + loaded = MessageFilterAgent.load_component(config) + assert loaded.name == "agent" + assert loaded._filter == wrapper._filter # pyright: ignore[reportPrivateUsage] + assert loaded._wrapped_agent.name == wrapper._wrapped_agent.name # pyright: ignore[reportPrivateUsage] + + # Run on_messages and validate filtering still works + messages = [ + TextMessage(source="user", content="u1"), + TextMessage(source="user", content="u2"), + TextMessage(source="user", content="u3"), + TextMessage(source="system", content="s1"), + TextMessage(source="system", content="s2"), + ] + await loaded.on_messages(messages, CancellationToken()) + received = loaded._wrapped_agent.received_messages # type: ignore[attr-defined] + assert {m.content for m in received} == {"u2", "u3", "s1"} # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType] + + +@pytest.mark.asyncio +async def test_message_filter_agent_in_digraph_group_chat(runtime: AgentRuntime | None) -> None: + inner_agent = _TestMessageFilterAgent("filtered") + filtered = MessageFilterAgent( + name="filtered", + wrapped_agent=inner_agent, + filter=MessageFilterConfig( + per_source=[ + PerSourceFilter(source="user", position="last", count=1), + ] + ), + ) + + graph = DiGraph( + nodes={ + "filtered": DiGraphNode(name="filtered", edges=[]), + } + ) + + team = GraphFlow( + participants=[filtered], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(3), + ) + + result = await team.run(task="only last user message matters") + assert result.stop_reason is not None + assert any(msg.source == "filtered" for msg in result.messages) + assert any(msg.content == "ACK" for msg in result.messages if msg.source == "filtered") # type: ignore[attr-defined,union-attr] + + +@pytest.mark.asyncio +async def test_message_filter_agent_loop_graph_visibility(runtime: AgentRuntime | None) -> None: + agent_a_inner = _TestMessageFilterAgent("A") + agent_a = MessageFilterAgent( + name="A", + wrapped_agent=agent_a_inner, + filter=MessageFilterConfig( + per_source=[ + PerSourceFilter(source="user", position="first", count=1), + PerSourceFilter(source="B", position="last", count=1), + ] + ), + ) + + from autogen_agentchat.agents import AssistantAgent + from autogen_ext.models.replay import ReplayChatCompletionClient + + model_client = ReplayChatCompletionClient(["loop", "loop", "exit"]) + agent_b_inner = AssistantAgent("B", model_client=model_client) + agent_b = MessageFilterAgent( + name="B", + wrapped_agent=agent_b_inner, + filter=MessageFilterConfig( + per_source=[ + PerSourceFilter(source="user", position="first", count=1), + PerSourceFilter(source="A", position="last", count=1), + PerSourceFilter(source="B", position="last", count=10), + ] + ), + ) + + agent_c_inner = _TestMessageFilterAgent("C") + agent_c = MessageFilterAgent( + name="C", + wrapped_agent=agent_c_inner, + filter=MessageFilterConfig( + per_source=[ + PerSourceFilter(source="user", position="first", count=1), + PerSourceFilter(source="B", position="last", count=1), + ] + ), + ) + + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode( + name="B", + edges=[ + DiGraphEdge(target="C", condition="exit"), + DiGraphEdge(target="A", condition="loop"), + ], + ), + "C": DiGraphNode(name="C", edges=[]), + }, + default_start_node="A", + ) + + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(20), + ) + + result = await team.run(task="Start") + assert result.stop_reason is not None + + # Check A received: 1 user + 2 from B + assert [m.source for m in agent_a_inner.received_messages].count("user") == 1 + assert [m.source for m in agent_a_inner.received_messages].count("B") == 2 + + # Check C received: 1 user + 1 from B + assert [m.source for m in agent_c_inner.received_messages].count("user") == 1 + assert [m.source for m in agent_c_inner.received_messages].count("B") == 1 + + # Check B received: 1 user + multiple from A + own messages + model_msgs = await agent_b_inner.model_context.get_messages() + sources = [m.source for m in model_msgs] # type: ignore[union-attr] + assert sources.count("user") == 1 # pyright: ignore[reportUnknownMemberType] + assert sources.count("A") >= 3 # pyright: ignore[reportUnknownMemberType] + assert sources.count("B") >= 2 # pyright: ignore[reportUnknownMemberType] + + +# Test Graph Builder +def test_add_node() -> None: + client = ReplayChatCompletionClient(["response"]) + agent = AssistantAgent("A", model_client=client) + builder = DiGraphBuilder() + builder.add_node(agent) + + assert "A" in builder.nodes + assert "A" in builder.agents + assert builder.nodes["A"].activation == "all" + + +def test_add_edge() -> None: + client = ReplayChatCompletionClient(["1", "2"]) + a = AssistantAgent("A", model_client=client) + b = AssistantAgent("B", model_client=client) + + builder = DiGraphBuilder() + builder.add_node(a).add_node(b) + builder.add_edge(a, b) + + assert builder.nodes["A"].edges[0].target == "B" + assert builder.nodes["A"].edges[0].condition is None + + +def test_add_conditional_edges() -> None: + client = ReplayChatCompletionClient(["1", "2"]) + a = AssistantAgent("A", model_client=client) + b = AssistantAgent("B", model_client=client) + c = AssistantAgent("C", model_client=client) + + builder = DiGraphBuilder() + builder.add_node(a).add_node(b).add_node(c) + builder.add_conditional_edges(a, {"yes": b, "no": c}) + + edges = builder.nodes["A"].edges + assert len(edges) == 2 + + # Extract the condition strings to compare them + conditions = [e.condition for e in edges] + assert "yes" in conditions + assert "no" in conditions + + # Match edge targets with conditions + yes_edge = next(e for e in edges if e.condition == "yes") + no_edge = next(e for e in edges if e.condition == "no") + + assert yes_edge.target == "B" + assert no_edge.target == "C" + + +def test_set_entry_point() -> None: + client = ReplayChatCompletionClient(["ok"]) + a = AssistantAgent("A", model_client=client) + builder = DiGraphBuilder().add_node(a).set_entry_point(a) + graph = builder.build() + + assert graph.default_start_node == "A" + + +def test_build_graph_validation() -> None: + client = ReplayChatCompletionClient(["1", "2", "3"]) + a = AssistantAgent("A", model_client=client) + b = AssistantAgent("B", model_client=client) + c = AssistantAgent("C", model_client=client) + + builder = DiGraphBuilder() + builder.add_node(a).add_node(b).add_node(c) + builder.add_edge("A", "B").add_edge("B", "C") + builder.set_entry_point("A") + graph = builder.build() + + assert isinstance(graph, DiGraph) + assert set(graph.nodes.keys()) == {"A", "B", "C"} + assert graph.get_start_nodes() == {"A"} + assert graph.get_leaf_nodes() == {"C"} + + +def test_build_fan_out() -> None: + client = ReplayChatCompletionClient(["hi"] * 3) + a = AssistantAgent("A", model_client=client) + b = AssistantAgent("B", model_client=client) + c = AssistantAgent("C", model_client=client) + + builder = DiGraphBuilder() + builder.add_node(a).add_node(b).add_node(c) + builder.add_edge(a, b).add_edge(a, c) + builder.set_entry_point(a) + graph = builder.build() + + assert graph.get_start_nodes() == {"A"} + assert graph.get_leaf_nodes() == {"B", "C"} + + +def test_build_parallel_join() -> None: + client = ReplayChatCompletionClient(["go"] * 3) + a = AssistantAgent("A", model_client=client) + b = AssistantAgent("B", model_client=client) + c = AssistantAgent("C", model_client=client) + + builder = DiGraphBuilder() + builder.add_node(a).add_node(b).add_node(c, activation="all") + builder.add_edge(a, c).add_edge(b, c) + builder.set_entry_point(a) + builder.add_edge(b, c) + builder.nodes["B"] = DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]) + graph = builder.build() + + assert graph.nodes["C"].activation == "all" + assert graph.get_leaf_nodes() == {"C"} + + +def test_build_conditional_loop() -> None: + client = ReplayChatCompletionClient(["loop", "loop", "exit"]) + a = AssistantAgent("A", model_client=client) + b = AssistantAgent("B", model_client=client) + c = AssistantAgent("C", model_client=client) + + builder = DiGraphBuilder() + builder.add_node(a).add_node(b).add_node(c) + builder.add_edge(a, b) + builder.add_conditional_edges(b, {"loop": a, "exit": c}) + builder.set_entry_point(a) + graph = builder.build() + + # Check that edges have the right conditions and targets + edges = graph.nodes["B"].edges + assert len(edges) == 2 + + # Find edges by their conditions + loop_edge = next(e for e in edges if e.condition == "loop") + exit_edge = next(e for e in edges if e.condition == "exit") + + assert loop_edge.target == "A" + assert exit_edge.target == "C" + assert graph.has_cycles_with_exit() + + +@pytest.mark.asyncio +async def test_graph_builder_sequential_execution(runtime: AgentRuntime | None) -> None: + a = _EchoAgent("A", description="Echo A") + b = _EchoAgent("B", description="Echo B") + c = _EchoAgent("C", description="Echo C") + + builder = DiGraphBuilder() + builder.add_node(a).add_node(b).add_node(c) + builder.add_edge(a, b).add_edge(b, c) + + team = GraphFlow( + participants=builder.get_participants(), + graph=builder.build(), + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + + result = await team.run(task="Start") + assert [m.source for m in result.messages[1:]] == ["A", "B", "C"] + assert result.stop_reason is not None + + +@pytest.mark.asyncio +async def test_graph_builder_fan_out(runtime: AgentRuntime | None) -> None: + a = _EchoAgent("A", description="Echo A") + b = _EchoAgent("B", description="Echo B") + c = _EchoAgent("C", description="Echo C") + + builder = DiGraphBuilder() + builder.add_node(a).add_node(b).add_node(c) + builder.add_edge(a, b).add_edge(a, c) + + team = GraphFlow( + participants=builder.get_participants(), + graph=builder.build(), + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + + result = await team.run(task="Start") + sources = [m.source for m in result.messages if isinstance(m, TextMessage)] + assert set(sources[1:]) == {"A", "B", "C"} + assert result.stop_reason is not None + + +@pytest.mark.asyncio +async def test_graph_builder_conditional_execution(runtime: AgentRuntime | None) -> None: + a = _EchoAgent("A", description="Echo A") + b = _EchoAgent("B", description="Echo B") + c = _EchoAgent("C", description="Echo C") + + builder = DiGraphBuilder() + builder.add_node(a).add_node(b).add_node(c) + builder.add_conditional_edges(a, {"yes": b, "no": c}) + + team = GraphFlow( + participants=builder.get_participants(), + graph=builder.build(), + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + + # Input "no" should trigger the edge to C + result = await team.run(task="no") + sources = [m.source for m in result.messages] + assert "C" in sources + assert result.stop_reason is not None + + +@pytest.mark.asyncio +async def test_digraph_group_chat_callable_condition(runtime: AgentRuntime | None) -> None: + """Test that string conditions work correctly in edge transitions.""" + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + + graph = DiGraph( + nodes={ + "A": DiGraphNode( + name="A", + edges=[ + # Will go to B if "long" is in message + DiGraphEdge(target="B", condition="long"), + # Will go to C if "short" is in message + DiGraphEdge(target="C", condition="short"), + ], + ), + "B": DiGraphNode(name="B", edges=[]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + + # Test with a message containing "long" - should go to B + result = await team.run(task="This is a long message") + assert result.messages[2].source == "B" + + # Reset for next test + await team.reset() + + # Test with a message containing "short" - should go to C + result = await team.run(task="This is a short message") + assert result.messages[2].source == "C" + + +@pytest.mark.asyncio +async def test_graph_flow_serialize_deserialize() -> None: + client_a = ReplayChatCompletionClient(list(map(str, range(10)))) + client_b = ReplayChatCompletionClient(list(map(str, range(10)))) + a = AssistantAgent("A", model_client=client_a) + b = AssistantAgent("B", model_client=client_b) + + builder = DiGraphBuilder() + builder.add_node(a).add_node(b) + builder.add_edge(a, b) + builder.set_entry_point(a) + + team = GraphFlow( + participants=builder.get_participants(), + graph=builder.build(), + runtime=None, + ) + + serialized = team.dump_component() + deserialized_team = GraphFlow.load_component(serialized) + serialized_deserialized = deserialized_team.dump_component() + + results = await team.run(task="Start") + de_results = await deserialized_team.run(task="Start") + + assert serialized == serialized_deserialized + assert compare_task_results(results, de_results) + assert results.stop_reason is not None + assert results.stop_reason == de_results.stop_reason + assert compare_message_lists(results.messages, de_results.messages) + assert isinstance(results.messages[0], TextMessage) + assert results.messages[0].source == "user" + assert results.messages[0].content == "Start" + assert isinstance(results.messages[1], TextMessage) + assert results.messages[1].source == "A" + assert results.messages[1].content == "0" + assert isinstance(results.messages[2], TextMessage) + assert results.messages[2].source == "B" + assert results.messages[2].content == "0" + # No stop agent message should appear in the conversation + assert all(not isinstance(m, StopMessage) for m in results.messages) + assert results.stop_reason is not None + + +@pytest.mark.asyncio +async def test_graph_flow_stateful_pause_and_resume_with_termination() -> None: + client_a = ReplayChatCompletionClient(["A1", "A2"]) + client_b = ReplayChatCompletionClient(["B1"]) + + a = AssistantAgent("A", model_client=client_a) + b = AssistantAgent("B", model_client=client_b) + + builder = DiGraphBuilder() + builder.add_node(a).add_node(b) + builder.add_edge(a, b) + builder.set_entry_point(a) + + team = GraphFlow( + participants=builder.get_participants(), + graph=builder.build(), + runtime=None, + termination_condition=SourceMatchTermination(sources=["A"]), + ) + + result = await team.run(task="Start") + assert len(result.messages) == 2 + assert result.messages[0].source == "user" + assert result.messages[1].source == "A" + assert result.stop_reason is not None and result.stop_reason == "'A' answered" + + # Export state. + state = await team.save_state() + + # Load state into a new team. + new_team = GraphFlow( + participants=builder.get_participants(), + graph=builder.build(), + runtime=None, + ) + await new_team.load_state(state) + + # Resume. + result = await new_team.run() + assert len(result.messages) == 1 + assert result.messages[0].source == "B" + + +@pytest.mark.asyncio +async def test_builder_with_lambda_condition(runtime: AgentRuntime | None) -> None: + """Test that DiGraphBuilder supports string conditions.""" + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + + builder = DiGraphBuilder() + builder.add_node(agent_a).add_node(agent_b).add_node(agent_c) + + # Using callable conditions + builder.add_edge(agent_a, agent_b, lambda msg: "even" in msg.to_model_text()) + builder.add_edge(agent_a, agent_c, lambda msg: "odd" in msg.to_model_text()) + + team = GraphFlow( + participants=builder.get_participants(), + graph=builder.build(), + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + + # Test with "even" in message - should go to B + result = await team.run(task="even length") + assert result.messages[2].source == "B" + + # Reset for next test + await team.reset() + + # Test with "odd" in message - should go to C + result = await team.run(task="odd message") + assert result.messages[2].source == "C" + + +@pytest.mark.asyncio +async def test_digraph_group_chat_multiple_task_execution(runtime: AgentRuntime | None) -> None: + """Test that GraphFlow can run multiple tasks sequentially after resetting execution state.""" + # Create agents A → B → C + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + + # Define graph A → B → C + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + + # Create team using Graph + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(5), + ) + + # Run the first task + result1: TaskResult = await team.run(task="First task") + + assert len(result1.messages) == 4 + assert isinstance(result1.messages[0], TextMessage) + assert result1.messages[0].source == "user" + assert result1.messages[0].content == "First task" + assert result1.messages[1].source == "A" + assert result1.messages[2].source == "B" + assert result1.messages[3].source == "C" + assert result1.stop_reason is not None + + # Run the second task - should work without explicit reset + result2: TaskResult = await team.run(task="Second task") + + assert len(result2.messages) == 4 + assert isinstance(result2.messages[0], TextMessage) + assert result2.messages[0].source == "user" + assert result2.messages[0].content == "Second task" + assert result2.messages[1].source == "A" + assert result2.messages[2].source == "B" + assert result2.messages[3].source == "C" + assert result2.stop_reason is not None + + # Verify agents were properly reset and executed again + assert agent_a.total_messages == 2 # Once for each task + assert agent_b.total_messages == 2 # Once for each task + assert agent_c.total_messages == 2 # Once for each task + + +@pytest.mark.asyncio +async def test_digraph_group_chat_resume_with_termination_condition(runtime: AgentRuntime | None) -> None: + """Test that GraphFlow can be resumed with the same execution state when a termination condition is reached.""" + # Create agents A → B → C + agent_a = _EchoAgent("A", description="Echo agent A") + agent_b = _EchoAgent("B", description="Echo agent B") + agent_c = _EchoAgent("C", description="Echo agent C") + + # Define graph A → B → C + graph = DiGraph( + nodes={ + "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]), + "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]), + "C": DiGraphNode(name="C", edges=[]), + } + ) + + # Create team with MaxMessageTermination that will stop before completion + team = GraphFlow( + participants=[agent_a, agent_b, agent_c], + graph=graph, + runtime=runtime, + termination_condition=MaxMessageTermination(3), # Stop after user + A + B + ) + + # Run the graph flow until termination condition is reached + result1: TaskResult = await team.run(task="Start execution") + + # Should have stopped at termination condition (user + A + B messages) + assert len(result1.messages) == 3 + assert result1.messages[0].source == "user" + assert result1.messages[1].source == "A" + assert result1.messages[2].source == "B" + assert result1.stop_reason is not None + + # Verify A and B ran, but C did not + assert agent_a.total_messages == 1 + assert agent_b.total_messages == 1 + assert agent_c.total_messages == 0 + + # Resume the graph flow with no task to continue where it left off + result2: TaskResult = await team.run() + + # Should continue and execute C, then complete without stop agent message + assert len(result2.messages) == 1 + assert result2.messages[0].source == "C" + assert result2.stop_reason is not None + + # Verify C now ran and the execution state was preserved + assert agent_a.total_messages == 1 # Still only ran once + assert agent_b.total_messages == 1 # Still only ran once + assert agent_c.total_messages == 1 # Now ran once diff --git a/python/packages/autogen-agentchat/tests/test_group_chat_nested.py b/python/packages/autogen-agentchat/tests/test_group_chat_nested.py new file mode 100644 index 000000000000..40301863ed20 --- /dev/null +++ b/python/packages/autogen-agentchat/tests/test_group_chat_nested.py @@ -0,0 +1,668 @@ +import logging +import tempfile +from collections.abc import AsyncGenerator + +import pytest +import pytest_asyncio +from autogen_agentchat import EVENT_LOGGER_NAME +from autogen_agentchat.agents import AssistantAgent, CodeExecutorAgent +from autogen_agentchat.base import TaskResult +from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination +from autogen_agentchat.messages import ( + BaseAgentEvent, + BaseChatMessage, + TextMessage, +) +from autogen_agentchat.teams import RoundRobinGroupChat, SelectorGroupChat, Swarm +from autogen_core import AgentRuntime, SingleThreadedAgentRuntime +from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor +from autogen_ext.models.replay import ReplayChatCompletionClient + +# Import test utilities from the main test file +from utils import FileLogHandler + +logger = logging.getLogger(EVENT_LOGGER_NAME) +logger.setLevel(logging.DEBUG) +logger.addHandler(FileLogHandler("test_group_chat_nested.log")) + + +@pytest_asyncio.fixture(params=["single_threaded", "embedded"]) # type: ignore +async def runtime(request: pytest.FixtureRequest) -> AsyncGenerator[AgentRuntime | None, None]: + if request.param == "single_threaded": + runtime = SingleThreadedAgentRuntime() + runtime.start() + yield runtime + await runtime.stop() + elif request.param == "embedded": + yield None + + +@pytest.mark.asyncio +async def test_round_robin_group_chat_nested_teams_run(runtime: AgentRuntime | None) -> None: + """Test RoundRobinGroupChat with nested teams using run method.""" + model_client = ReplayChatCompletionClient( + [ + 'Here is the program\n ```python\nprint("Hello, world!")\n```', + "TERMINATE", + "Good job", + "TERMINATE", + ], + ) + with tempfile.TemporaryDirectory() as temp_dir: + code_executor = LocalCommandLineCodeExecutor(work_dir=temp_dir) + assistant = AssistantAgent( + "assistant", + model_client=model_client, + description="An assistant agent that writes code.", + ) + code_executor_agent = CodeExecutorAgent("code_executor", code_executor=code_executor) + termination = TextMentionTermination("TERMINATE") + + # Create inner team (assistant + code executor) + inner_team = RoundRobinGroupChat( + participants=[assistant, code_executor_agent], + termination_condition=termination, + runtime=runtime, + ) + + # Create reviewer agent + reviewer = AssistantAgent( + "reviewer", + model_client=model_client, + description="A reviewer agent that reviews code.", + ) + + # Create outer team with nested inner team + outer_team = RoundRobinGroupChat( + participants=[inner_team, reviewer], + termination_condition=termination, + runtime=runtime, + ) + + result = await outer_team.run(task="Write a program that prints 'Hello, world!'") + + # Should have task message + inner team result + reviewer response + termination + assert len(result.messages) >= 4 + assert isinstance(result.messages[0], TextMessage) + assert result.messages[0].content == "Write a program that prints 'Hello, world!'" + assert result.stop_reason is not None and "TERMINATE" in result.stop_reason + + +@pytest.mark.asyncio +async def test_round_robin_group_chat_nested_teams_run_stream(runtime: AgentRuntime | None) -> None: + """Test RoundRobinGroupChat with nested teams using run_stream method.""" + model_client = ReplayChatCompletionClient( + [ + 'Here is the program\n ```python\nprint("Hello, world!")\n```', + "TERMINATE", + "Good job", + "TERMINATE", + ], + ) + with tempfile.TemporaryDirectory() as temp_dir: + code_executor = LocalCommandLineCodeExecutor(work_dir=temp_dir) + assistant = AssistantAgent( + "assistant", + model_client=model_client, + description="An assistant agent that writes code.", + ) + code_executor_agent = CodeExecutorAgent("code_executor", code_executor=code_executor) + termination = TextMentionTermination("TERMINATE") + + # Create inner team (assistant + code executor) + inner_team = RoundRobinGroupChat( + participants=[assistant, code_executor_agent], + termination_condition=termination, + runtime=runtime, + ) + + # Create reviewer agent + reviewer = AssistantAgent( + "reviewer", + model_client=model_client, + description="A reviewer agent that reviews code.", + ) + + # Create outer team with nested inner team + outer_team = RoundRobinGroupChat( + participants=[inner_team, reviewer], + termination_condition=termination, + runtime=runtime, + ) + + messages: list[BaseAgentEvent | BaseChatMessage] = [] + result = None + async for message in outer_team.run_stream(task="Write a program that prints 'Hello, world!'"): + if isinstance(message, TaskResult): + result = message + else: + messages.append(message) + + assert result is not None + assert len(result.messages) >= 4 + assert isinstance(result.messages[0], TextMessage) + assert result.messages[0].content == "Write a program that prints 'Hello, world!'" + assert result.stop_reason is not None and "TERMINATE" in result.stop_reason + + +@pytest.mark.asyncio +async def test_round_robin_group_chat_nested_teams_dump_load_component(runtime: AgentRuntime | None) -> None: + """Test RoundRobinGroupChat with nested teams dump_component and load_component.""" + model_client = ReplayChatCompletionClient(["Hello from agent1", "Hello from agent2", "Hello from agent3"]) + + # Create agents + agent1 = AssistantAgent("agent1", model_client=model_client, description="First agent") + agent2 = AssistantAgent("agent2", model_client=model_client, description="Second agent") + agent3 = AssistantAgent("agent3", model_client=model_client, description="Third agent") + termination = MaxMessageTermination(2) + + # Create inner team + inner_team = RoundRobinGroupChat( + participants=[agent1, agent2], + termination_condition=termination, + runtime=runtime, + name="InnerTeam", + description="Inner team description", + ) + + # Create outer team with nested inner team + outer_team = RoundRobinGroupChat( + participants=[inner_team, agent3], + termination_condition=termination, + runtime=runtime, + name="OuterTeam", + description="Outer team description", + ) + + # Test dump_component + config = outer_team.dump_component() + assert config.config["name"] == "OuterTeam" + assert config.config["description"] == "Outer team description" + assert len(config.config["participants"]) == 2 + + # First participant should be the inner team + inner_team_config = config.config["participants"][0]["config"] + assert inner_team_config["name"] == "InnerTeam" + assert inner_team_config["description"] == "Inner team description" + assert len(inner_team_config["participants"]) == 2 + + # Second participant should be agent3 + agent3_config = config.config["participants"][1]["config"] + assert agent3_config["name"] == "agent3" + + # Test load_component + loaded_team = RoundRobinGroupChat.load_component(config) + assert loaded_team.name == "OuterTeam" + assert loaded_team.description == "Outer team description" + assert len(loaded_team._participants) == 2 # type: ignore[reportPrivateUsage] + + # Verify the loaded team has the same structure + loaded_config = loaded_team.dump_component() + assert loaded_config == config + + +@pytest.mark.asyncio +async def test_round_robin_group_chat_nested_teams_save_load_state(runtime: AgentRuntime | None) -> None: + """Test RoundRobinGroupChat with nested teams save_state and load_state.""" + model_client = ReplayChatCompletionClient(["Hello from agent1", "Hello from agent2", "TERMINATE"]) + + # Create agents + agent1 = AssistantAgent("agent1", model_client=model_client, description="First agent") + agent2 = AssistantAgent("agent2", model_client=model_client, description="Second agent") + agent3 = AssistantAgent("agent3", model_client=model_client, description="Third agent") + termination = TextMentionTermination("TERMINATE") # Use TextMentionTermination + + # Create inner team + inner_team = RoundRobinGroupChat( + participants=[agent1, agent2], + termination_condition=termination, + runtime=runtime, + ) + + # Create outer team with nested inner team + outer_team1 = RoundRobinGroupChat( + participants=[inner_team, agent3], + termination_condition=termination, + runtime=runtime, + ) + + # Run the team to generate state + await outer_team1.run(task="Test message") + + # Save state + state = await outer_team1.save_state() + + # Create new agents and teams + agent4 = AssistantAgent("agent1", model_client=model_client, description="First agent") + agent5 = AssistantAgent("agent2", model_client=model_client, description="Second agent") + agent6 = AssistantAgent("agent3", model_client=model_client, description="Third agent") + + inner_team2 = RoundRobinGroupChat( + participants=[agent4, agent5], + termination_condition=termination, + runtime=runtime, + ) + + outer_team2 = RoundRobinGroupChat( + participants=[inner_team2, agent6], + termination_condition=termination, + runtime=runtime, + ) + + # Load state + await outer_team2.load_state(state) + + # Verify state was loaded correctly + state2 = await outer_team2.save_state() + assert state == state2 + + +@pytest.mark.asyncio +async def test_selector_group_chat_nested_teams_run(runtime: AgentRuntime | None) -> None: + """Test SelectorGroupChat with nested teams using run method.""" + model_client = ReplayChatCompletionClient( + [ + "InnerTeam", # Select inner team first + 'Here is the program\n ```python\nprint("Hello, world!")\n```', + "TERMINATE", + "agent3", # Select agent3 (reviewer) + "Good job", + "TERMINATE", + ], + ) + with tempfile.TemporaryDirectory() as temp_dir: + code_executor = LocalCommandLineCodeExecutor(work_dir=temp_dir) + assistant = AssistantAgent( + "assistant", + model_client=model_client, + description="An assistant agent that writes code.", + ) + code_executor_agent = CodeExecutorAgent("code_executor", code_executor=code_executor) + termination = TextMentionTermination("TERMINATE") + + # Create inner team (assistant + code executor) + inner_team = RoundRobinGroupChat( + participants=[assistant, code_executor_agent], + termination_condition=termination, + runtime=runtime, + name="InnerTeam", + description="Team that writes and executes code", + ) + + # Create reviewer agent + reviewer = AssistantAgent( + "agent3", + model_client=model_client, + description="A reviewer agent that reviews code.", + ) + + # Create outer team with nested inner team + outer_team = SelectorGroupChat( + participants=[inner_team, reviewer], + model_client=model_client, + termination_condition=termination, + runtime=runtime, + ) + + result = await outer_team.run(task="Write a program that prints 'Hello, world!'") + + # Should have task message + selector events + inner team result + reviewer response + assert len(result.messages) >= 4 + assert isinstance(result.messages[0], TextMessage) + assert result.messages[0].content == "Write a program that prints 'Hello, world!'" + assert result.stop_reason is not None and "TERMINATE" in result.stop_reason + + +@pytest.mark.asyncio +async def test_selector_group_chat_nested_teams_run_stream(runtime: AgentRuntime | None) -> None: + """Test SelectorGroupChat with nested teams using run_stream method.""" + model_client = ReplayChatCompletionClient( + [ + "InnerTeam", # Select inner team first + 'Here is the program\n ```python\nprint("Hello, world!")\n```', + "TERMINATE", + "agent3", # Select agent3 (reviewer) + "Good job", + "TERMINATE", + ], + ) + with tempfile.TemporaryDirectory() as temp_dir: + code_executor = LocalCommandLineCodeExecutor(work_dir=temp_dir) + assistant = AssistantAgent( + "assistant", + model_client=model_client, + description="An assistant agent that writes code.", + ) + code_executor_agent = CodeExecutorAgent("code_executor", code_executor=code_executor) + termination = TextMentionTermination("TERMINATE") + + # Create inner team (assistant + code executor) + inner_team = RoundRobinGroupChat( + participants=[assistant, code_executor_agent], + termination_condition=termination, + runtime=runtime, + name="InnerTeam", + description="Team that writes and executes code", + ) + + # Create reviewer agent + reviewer = AssistantAgent( + "agent3", + model_client=model_client, + description="A reviewer agent that reviews code.", + ) + + # Create outer team with nested inner team + outer_team = SelectorGroupChat( + participants=[inner_team, reviewer], + model_client=model_client, + termination_condition=termination, + runtime=runtime, + ) + + messages: list[BaseAgentEvent | BaseChatMessage] = [] + result = None + async for message in outer_team.run_stream(task="Write a program that prints 'Hello, world!'"): + if isinstance(message, TaskResult): + result = message + else: + messages.append(message) + + assert result is not None + assert len(result.messages) >= 4 + assert isinstance(result.messages[0], TextMessage) + assert result.messages[0].content == "Write a program that prints 'Hello, world!'" + assert result.stop_reason is not None and "TERMINATE" in result.stop_reason + + +@pytest.mark.asyncio +async def test_selector_group_chat_nested_teams_dump_load_component(runtime: AgentRuntime | None) -> None: + """Test SelectorGroupChat with nested teams dump_component and load_component.""" + model_client = ReplayChatCompletionClient(["agent1", "Hello from agent1", "agent3", "Hello from agent3"]) + + # Create agents + agent1 = AssistantAgent("agent1", model_client=model_client, description="First agent") + agent2 = AssistantAgent("agent2", model_client=model_client, description="Second agent") + agent3 = AssistantAgent("agent3", model_client=model_client, description="Third agent") + termination = MaxMessageTermination(2) + + # Create inner team + inner_team = RoundRobinGroupChat( + participants=[agent1, agent2], + termination_condition=termination, + runtime=runtime, + name="InnerTeam", + description="Inner team description", + ) + + # Create outer team with nested inner team + outer_team = SelectorGroupChat( + participants=[inner_team, agent3], + model_client=model_client, + termination_condition=termination, + runtime=runtime, + name="OuterTeam", + description="Outer team description", + ) + + # Test dump_component + config = outer_team.dump_component() + assert config.config["name"] == "OuterTeam" + assert config.config["description"] == "Outer team description" + assert len(config.config["participants"]) == 2 + + # First participant should be the inner team + inner_team_config = config.config["participants"][0]["config"] + assert inner_team_config["name"] == "InnerTeam" + assert inner_team_config["description"] == "Inner team description" + assert len(inner_team_config["participants"]) == 2 + + # Second participant should be agent3 + agent3_config = config.config["participants"][1]["config"] + assert agent3_config["name"] == "agent3" + + # Test load_component + loaded_team = SelectorGroupChat.load_component(config) + assert loaded_team.name == "OuterTeam" + assert loaded_team.description == "Outer team description" + assert len(loaded_team._participants) == 2 # type: ignore[reportPrivateUsage] + + # Verify the loaded team has the same structure + loaded_config = loaded_team.dump_component() + assert loaded_config == config + + +@pytest.mark.asyncio +async def test_selector_group_chat_nested_teams_save_load_state(runtime: AgentRuntime | None) -> None: + """Test SelectorGroupChat with nested teams save_state and load_state.""" + model_client = ReplayChatCompletionClient(["InnerTeam", "Hello from inner team", "agent3", "TERMINATE"]) + + # Create agents + agent1 = AssistantAgent("agent1", model_client=model_client, description="First agent") + agent2 = AssistantAgent("agent2", model_client=model_client, description="Second agent") + agent3 = AssistantAgent("agent3", model_client=model_client, description="Third agent") + termination = TextMentionTermination("TERMINATE") + + # Create inner team + inner_team = RoundRobinGroupChat( + participants=[agent1, agent2], + termination_condition=termination, + runtime=runtime, + name="InnerTeam", + ) + + # Create outer team with nested inner team + outer_team1 = SelectorGroupChat( + participants=[inner_team, agent3], + model_client=model_client, + termination_condition=termination, + runtime=runtime, + ) + + # Run the team to generate state + await outer_team1.run(task="Test message") + + # Save state + state = await outer_team1.save_state() + + # Create new agents and teams + agent4 = AssistantAgent("agent1", model_client=model_client, description="First agent") + agent5 = AssistantAgent("agent2", model_client=model_client, description="Second agent") + agent6 = AssistantAgent("agent3", model_client=model_client, description="Third agent") + + inner_team2 = RoundRobinGroupChat( + participants=[agent4, agent5], + termination_condition=termination, + runtime=runtime, + name="InnerTeam", + ) + + outer_team2 = SelectorGroupChat( + participants=[inner_team2, agent6], + model_client=model_client, + termination_condition=termination, + runtime=runtime, + ) + + # Load state + await outer_team2.load_state(state) + + # Verify state was loaded correctly + state2 = await outer_team2.save_state() + assert state == state2 + + +@pytest.mark.asyncio +async def test_swarm_doesnt_support_nested_teams() -> None: + """Test that Swarm raises TypeError when provided with nested teams.""" + model_client = ReplayChatCompletionClient(["Hello", "TERMINATE"]) + + # Create agents + agent1 = AssistantAgent("agent1", model_client=model_client, description="First agent") + agent2 = AssistantAgent("agent2", model_client=model_client, description="Second agent") + agent3 = AssistantAgent("agent3", model_client=model_client, description="Third agent") + termination = TextMentionTermination("TERMINATE") + + # Create inner team + inner_team = RoundRobinGroupChat( + participants=[agent1, agent2], + termination_condition=termination, + ) + + # Verify that Swarm raises TypeError when trying to use a team as participant + with pytest.raises(TypeError, match="Participant .* must be a ChatAgent"): + Swarm( + participants=[inner_team, agent3], # type: ignore + termination_condition=termination, + ) + + +@pytest.mark.asyncio +async def test_round_robin_deeply_nested_teams(runtime: AgentRuntime | None) -> None: + """Test RoundRobinGroupChat with deeply nested teams (3 levels).""" + model_client = ReplayChatCompletionClient( + [ + "Hello from agent1", + "TERMINATE from agent2", + "World from agent3", + "Hello from agent1", + "Hello from agent2", + "TERMINATE from agent1", + "TERMINATE from agent3", + "Review from agent4", + "TERMINATE from agent2", + "TERMINATE from agent3", + "TERMINATE from agent4", + ] + ) + + # Create agents + agent1 = AssistantAgent("agent1", model_client=model_client, description="First agent") + agent2 = AssistantAgent("agent2", model_client=model_client, description="Second agent") + agent3 = AssistantAgent("agent3", model_client=model_client, description="Third agent") + agent4 = AssistantAgent("agent4", model_client=model_client, description="Fourth agent") + + # Create innermost team (level 1) + innermost_team = RoundRobinGroupChat( + participants=[agent1, agent2], + termination_condition=TextMentionTermination("TERMINATE", sources=["agent1", "agent2"]), + runtime=runtime, + name="InnermostTeam", + ) + + # Create middle team (level 2) + middle_team = RoundRobinGroupChat( + participants=[innermost_team, agent3], + termination_condition=TextMentionTermination("TERMINATE", sources=["agent3"]), + runtime=runtime, + name="MiddleTeam", + ) + + # Create outermost team (level 3) + outermost_team = RoundRobinGroupChat( + participants=[middle_team, agent4], + termination_condition=TextMentionTermination("TERMINATE", sources=["agent4"]), + runtime=runtime, + name="OutermostTeam", + ) + + result: TaskResult | None = None + async for msg in outermost_team.run_stream(task="Test deep nesting"): + if isinstance(msg, TaskResult): + result = msg + assert result is not None + # Should have task message + responses from each level + assert len(result.messages) == 12 + assert isinstance(result.messages[0], TextMessage) + assert result.messages[0].content == "Test deep nesting" + assert result.stop_reason is not None and "TERMINATE" in result.stop_reason + + # Test component serialization of deeply nested structure + config = outermost_team.dump_component() + loaded_team = RoundRobinGroupChat.load_component(config) + assert loaded_team.name == "OutermostTeam" + + # Verify nested structure is preserved + loaded_config = loaded_team.dump_component() + assert loaded_config == config + + +@pytest.mark.asyncio +async def test_selector_deeply_nested_teams(runtime: AgentRuntime | None) -> None: + """Test SelectorGroupChat with deeply nested teams (3 levels).""" + model_client_inner = ReplayChatCompletionClient( + [ + "Hello from innermost agent 1", + "Hello from innermost agent 2", + "TERMINATE from innermost agent 1", + ] + ) + model_client_middle = ReplayChatCompletionClient( + [ + "InnermostTeam", # Select innermost team + "TERMINATE from agent3", + ] + ) + model_client_outter = ReplayChatCompletionClient( + [ + "MiddleTeam", # Select middle team + "agent4", # Select agent4 + "Hello from outermost agent 4", + "agent4", # Select agent4 again + "TERMINATE from agent4", + ] + ) + + # Create agents + agent1 = AssistantAgent("agent1", model_client=model_client_inner, description="First agent") + agent2 = AssistantAgent("agent2", model_client=model_client_inner, description="Second agent") + agent3 = AssistantAgent("agent3", model_client=model_client_middle, description="Third agent") + agent4 = AssistantAgent("agent4", model_client=model_client_outter, description="Fourth agent") + + # Create innermost team (level 1) - RoundRobin for simplicity + innermost_team = RoundRobinGroupChat( + participants=[agent1, agent2], + termination_condition=TextMentionTermination("TERMINATE", sources=["agent1", "agent2"]), + runtime=runtime, + name="InnermostTeam", + ) + + # Create middle team (level 2) - Selector + middle_team = SelectorGroupChat( + participants=[innermost_team, agent3], + model_client=model_client_middle, + termination_condition=TextMentionTermination("TERMINATE", sources=["agent3"]), + runtime=runtime, + name="MiddleTeam", + ) + + # Create outermost team (level 3) - Selector + outermost_team = SelectorGroupChat( + participants=[middle_team, agent4], + model_client=model_client_outter, + termination_condition=TextMentionTermination("TERMINATE", sources=["agent4"]), + runtime=runtime, + name="OutermostTeam", + allow_repeated_speaker=True, + ) + + result: TaskResult | None = None + async for msg in outermost_team.run_stream(task="Test deep nesting"): + if isinstance(msg, TaskResult): + result = msg + assert result is not None + + # Should have task message + selector events + responses from each level + assert len(result.messages) == 7 + assert isinstance(result.messages[0], TextMessage) + assert result.messages[0].content == "Test deep nesting" + assert result.stop_reason is not None and "TERMINATE" in result.stop_reason + + # Test component serialization of deeply nested structure + config = outermost_team.dump_component() + loaded_team = SelectorGroupChat.load_component(config) + assert loaded_team.name == "OutermostTeam" + + # Verify nested structure is preserved + loaded_config = loaded_team.dump_component() + assert loaded_config == config diff --git a/python/packages/autogen-agentchat/tests/test_group_chat_pause_resume.py b/python/packages/autogen-agentchat/tests/test_group_chat_pause_resume.py index ee87f6f4a25b..e26c7262d66d 100644 --- a/python/packages/autogen-agentchat/tests/test_group_chat_pause_resume.py +++ b/python/packages/autogen-agentchat/tests/test_group_chat_pause_resume.py @@ -5,7 +5,7 @@ import pytest_asyncio from autogen_agentchat.agents import BaseChatAgent from autogen_agentchat.base import Response -from autogen_agentchat.messages import ChatMessage, TextMessage +from autogen_agentchat.messages import BaseChatMessage, TextMessage from autogen_agentchat.teams import RoundRobinGroupChat from autogen_core import AgentRuntime, CancellationToken, SingleThreadedAgentRuntime @@ -20,10 +20,10 @@ def __init__(self, name: str, description: str) -> None: self.counter = 0 @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: return [TextMessage] - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: assert not self._is_paused, "Agent is paused" async def _process() -> None: diff --git a/python/packages/autogen-agentchat/tests/test_magentic_one_group_chat.py b/python/packages/autogen-agentchat/tests/test_magentic_one_group_chat.py index 34228ba039f7..d217d54763c7 100644 --- a/python/packages/autogen-agentchat/tests/test_magentic_one_group_chat.py +++ b/python/packages/autogen-agentchat/tests/test_magentic_one_group_chat.py @@ -11,7 +11,7 @@ ) from autogen_agentchat.base import Response from autogen_agentchat.messages import ( - ChatMessage, + BaseChatMessage, TextMessage, ) from autogen_agentchat.teams import ( @@ -34,14 +34,14 @@ def __init__(self, name: str, description: str) -> None: self._total_messages = 0 @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: return (TextMessage,) @property def total_messages(self) -> int: return self._total_messages - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: if len(messages) > 0: assert isinstance(messages[0], TextMessage) self._last_message = messages[0].content @@ -134,8 +134,8 @@ async def test_magentic_one_group_chat_basic(runtime: AgentRuntime | None) -> No ) result = await team.run(task="Write a program that prints 'Hello, world!'") assert len(result.messages) == 5 - assert result.messages[2].content == "Continue task" - assert result.messages[4].content == "print('Hello, world!')" + assert result.messages[2].to_text() == "Continue task" + assert result.messages[4].to_text() == "print('Hello, world!')" assert result.stop_reason is not None and result.stop_reason == "Because" # Test save and load. @@ -214,8 +214,8 @@ async def test_magentic_one_group_chat_with_stalls(runtime: AgentRuntime | None) ) result = await team.run(task="Write a program that prints 'Hello, world!'") assert len(result.messages) == 6 - assert isinstance(result.messages[1].content, str) + assert isinstance(result.messages[1], TextMessage) assert result.messages[1].content.startswith("\nWe are working to address the following user request:") - assert isinstance(result.messages[4].content, str) + assert isinstance(result.messages[4], TextMessage) assert result.messages[4].content.startswith("\nWe are working to address the following user request:") assert result.stop_reason is not None and result.stop_reason == "test" diff --git a/python/packages/autogen-agentchat/tests/test_messages.py b/python/packages/autogen-agentchat/tests/test_messages.py new file mode 100644 index 000000000000..9718af3de705 --- /dev/null +++ b/python/packages/autogen-agentchat/tests/test_messages.py @@ -0,0 +1,368 @@ +import json +import uuid +from datetime import datetime, timezone +from typing import List + +import pytest +from autogen_agentchat.messages import ( + AgentEvent, + ChatMessage, + HandoffMessage, + MessageFactory, + ModelClientStreamingChunkEvent, + MultiModalMessage, + StopMessage, + StructuredMessage, + StructuredMessageFactory, + TextMessage, + ToolCallExecutionEvent, + ToolCallRequestEvent, +) +from autogen_core import FunctionCall +from autogen_core.models import FunctionExecutionResult +from pydantic import BaseModel + + +class TestContent(BaseModel): + """Test content model.""" + + field1: str + field2: int + + +def test_structured_message() -> None: + # Create a structured message with the test content + message = StructuredMessage[TestContent]( + source="test_agent", + content=TestContent(field1="test", field2=42), + ) + + # Check that the message type is correct + assert message.type == "StructuredMessage[TestContent]" # type: ignore[comparison-overlap] + + # Check that the content is of the correct type + assert isinstance(message.content, TestContent) + + # Check that the content fields are set correctly + assert message.content.field1 == "test" + assert message.content.field2 == 42 + + # Check that model_dump works correctly + dumped_message = message.model_dump() + assert dumped_message["source"] == "test_agent" + assert dumped_message["content"]["field1"] == "test" + assert dumped_message["content"]["field2"] == 42 + assert dumped_message["type"] == "StructuredMessage[TestContent]" + + +def test_structured_message_component() -> None: + # Create a structured message with the test content + format_string = "this is a string {field1} and this is an int {field2}" + s_m = StructuredMessageFactory(input_model=TestContent, format_string=format_string) + config = s_m.dump_component() + s_m_dyn = StructuredMessageFactory.load_component(config) + message = s_m_dyn.StructuredMessage( + source="test_agent", content=s_m_dyn.ContentModel(field1="test", field2=42), format_string=s_m_dyn.format_string + ) + + assert isinstance(message.content, s_m_dyn.ContentModel) + assert not isinstance(message.content, TestContent) + assert message.content.field1 == "test" # type: ignore[attr-defined] + assert message.content.field2 == 42 # type: ignore[attr-defined] + + dumped_message = message.model_dump() + assert dumped_message["source"] == "test_agent" + assert dumped_message["content"]["field1"] == "test" + assert dumped_message["content"]["field2"] == 42 + assert message.to_model_text() == format_string.format(field1="test", field2=42) + + +def test_message_factory() -> None: + factory = MessageFactory() + + # Text message data + text_data = { + "type": "TextMessage", + "source": "test_agent", + "content": "Hello, world!", + } + + # Create a TextMessage instance + text_message = factory.create(text_data) + assert isinstance(text_message, TextMessage) + assert text_message.source == "test_agent" + assert text_message.content == "Hello, world!" + assert text_message.type == "TextMessage" # type: ignore[comparison-overlap] + + # Handoff message data + handoff_data = { + "type": "HandoffMessage", + "source": "test_agent", + "content": "handoff to another agent", + "target": "target_agent", + } + + # Create a HandoffMessage instance + handoff_message = factory.create(handoff_data) + assert isinstance(handoff_message, HandoffMessage) + assert handoff_message.source == "test_agent" + assert handoff_message.content == "handoff to another agent" + assert handoff_message.target == "target_agent" + assert handoff_message.type == "HandoffMessage" # type: ignore[comparison-overlap] + + # Structured message data + structured_data = { + "type": "StructuredMessage[TestContent]", + "source": "test_agent", + "content": { + "field1": "test", + "field2": 42, + }, + } + # Create a StructuredMessage instance -- this will fail because the type + # is not registered in the factory. + with pytest.raises(ValueError): + structured_message = factory.create(structured_data) + # Register the StructuredMessage type in the factory + factory.register(StructuredMessage[TestContent]) + # Create a StructuredMessage instance + structured_message = factory.create(structured_data) + assert isinstance(structured_message, StructuredMessage) + assert isinstance(structured_message.content, TestContent) # type: ignore[reportUnkownMemberType] + assert structured_message.source == "test_agent" + assert structured_message.content.field1 == "test" + assert structured_message.content.field2 == 42 + assert structured_message.type == "StructuredMessage[TestContent]" # type: ignore[comparison-overlap] + + sm_factory = StructuredMessageFactory(input_model=TestContent, format_string=None, content_model_name="TestContent") + config = sm_factory.dump_component() + config.config["content_model_name"] = "DynamicTestContent" + sm_factory_dynamic = StructuredMessageFactory.load_component(config) + + factory.register(sm_factory_dynamic.StructuredMessage) + msg = sm_factory_dynamic.StructuredMessage( + content=sm_factory_dynamic.ContentModel(field1="static", field2=123), source="static_agent" + ) + restored = factory.create(msg.dump()) + assert isinstance(restored, StructuredMessage) + assert isinstance(restored.content, sm_factory_dynamic.ContentModel) # type: ignore[reportUnkownMemberType] + assert restored.source == "static_agent" + assert restored.content.field1 == "static" # type: ignore[attr-defined] + assert restored.content.field2 == 123 # type: ignore[attr-defined] + + +class TestContainer(BaseModel): + chat_messages: List[ChatMessage] + agent_events: List[AgentEvent] + + +def test_union_types() -> None: + # Create a few messages. + chat_messages: List[ChatMessage] = [ + TextMessage(source="user", content="Hello!"), + MultiModalMessage(source="user", content=["Hello!", "World!"]), + HandoffMessage(source="user", content="handoff to another agent", target="target_agent"), + StopMessage(source="user", content="stop"), + ] + + # Create a few agent events. + agent_events: List[AgentEvent] = [ + ModelClientStreamingChunkEvent(source="user", content="Hello!"), + ToolCallRequestEvent( + content=[ + FunctionCall(id="1", name="test_function", arguments=json.dumps({"arg1": "value1", "arg2": "value2"})) + ], + source="user", + ), + ToolCallExecutionEvent( + content=[FunctionExecutionResult(call_id="1", content="result", name="test")], source="user" + ), + ] + + # Create a container with the messages. + container = TestContainer(chat_messages=chat_messages, agent_events=agent_events) + + # Dump the container to JSON. + data = container.model_dump() + + # Load the container from JSON. + loaded_container = TestContainer.model_validate(data) + assert loaded_container.chat_messages == chat_messages + assert loaded_container.agent_events == agent_events + + +def test_message_id_field() -> None: + """Test that messages have unique ID fields automatically generated.""" + # Test BaseChatMessage subclass (TextMessage) + message1 = TextMessage(source="test_agent", content="Hello, world!") + message2 = TextMessage(source="test_agent", content="Hello, world!") + + # Check that IDs are present and unique + assert hasattr(message1, "id") + assert hasattr(message2, "id") + assert message1.id != message2.id + assert isinstance(message1.id, str) + assert isinstance(message2.id, str) + + # Check that IDs are valid UUIDs + try: + uuid.UUID(message1.id) + uuid.UUID(message2.id) + except ValueError: + pytest.fail("Generated IDs are not valid UUIDs") + + # Test BaseAgentEvent subclass (ModelClientStreamingChunkEvent) + event1 = ModelClientStreamingChunkEvent(source="test_agent", content="chunk1") + event2 = ModelClientStreamingChunkEvent(source="test_agent", content="chunk2") + + # Check that IDs are present and unique + assert hasattr(event1, "id") + assert hasattr(event2, "id") + assert event1.id != event2.id + assert isinstance(event1.id, str) + assert isinstance(event2.id, str) + + # Check that IDs are valid UUIDs + try: + uuid.UUID(event1.id) + uuid.UUID(event2.id) + except ValueError: + pytest.fail("Generated IDs are not valid UUIDs") + + +def test_custom_message_id() -> None: + """Test that custom IDs can be provided.""" + custom_id = "custom-message-id-123" + message = TextMessage(id=custom_id, source="test_agent", content="Hello, world!") + + assert message.id == custom_id + + custom_event_id = "custom-event-id-456" + event = ModelClientStreamingChunkEvent(id=custom_event_id, source="test_agent", content="chunk") + + assert event.id == custom_event_id + + +def test_streaming_chunk_full_message_id() -> None: + """Test the full_message_id field in ModelClientStreamingChunkEvent.""" + # Test without full_message_id + chunk1 = ModelClientStreamingChunkEvent(source="test_agent", content="chunk1") + assert chunk1.full_message_id is None + + # Test with full_message_id + full_msg_id = "full-message-123" + chunk2 = ModelClientStreamingChunkEvent(source="test_agent", content="chunk2", full_message_id=full_msg_id) + assert chunk2.full_message_id == full_msg_id + + # Test that chunk has its own ID separate from full_message_id + assert chunk2.id != chunk2.full_message_id + assert isinstance(chunk2.id, str) + + # Verify chunk ID is a valid UUID + try: + uuid.UUID(chunk2.id) + except ValueError: + pytest.fail("Chunk ID is not a valid UUID") + + +def test_message_serialization_with_id() -> None: + """Test that messages with IDs serialize and deserialize correctly.""" + # Create a message with auto-generated ID + original_message = TextMessage(source="test_agent", content="Hello, world!") + original_id = original_message.id + + # Serialize to dict + message_data = original_message.model_dump() + assert "id" in message_data + assert message_data["id"] == original_id + + # Deserialize from dict + restored_message = TextMessage.model_validate(message_data) + assert restored_message.id == original_id + assert restored_message.source == "test_agent" + assert restored_message.content == "Hello, world!" + + # Test with streaming chunk event + original_chunk = ModelClientStreamingChunkEvent( + source="test_agent", content="chunk", full_message_id="full-msg-123" + ) + original_chunk_id = original_chunk.id + + # Serialize to dict + chunk_data = original_chunk.model_dump() + assert "id" in chunk_data + assert "full_message_id" in chunk_data + assert chunk_data["id"] == original_chunk_id + assert chunk_data["full_message_id"] == "full-msg-123" + + # Deserialize from dict + restored_chunk = ModelClientStreamingChunkEvent.model_validate(chunk_data) + assert restored_chunk.id == original_chunk_id + assert restored_chunk.full_message_id == "full-msg-123" + assert restored_chunk.content == "chunk" + + +def test_datetime_serialization_in_messages() -> None: + """Test that datetime objects in messages are properly serialized to JSON-compatible format. + + This test validates the fix for issue #6793 where datetime objects in message + created_at fields caused JSON serialization errors when saving team state. + """ + # Create a specific datetime for testing + test_datetime = datetime(2023, 12, 25, 10, 30, 45, 123456, timezone.utc) + + # Test BaseChatMessage subclass with datetime + chat_message = TextMessage(source="test_agent", content="Hello, world!", created_at=test_datetime) + + # Test that dump() returns JSON-serializable data + chat_message_data = chat_message.dump() + + # Verify that the datetime is converted to a string in ISO format + assert isinstance(chat_message_data["created_at"], str) + # Pydantic JSON mode converts UTC timezone to 'Z' format instead of '+00:00' + expected_iso = test_datetime.isoformat().replace("+00:00", "Z") + assert chat_message_data["created_at"] == expected_iso + + # Verify that the dumped data is JSON serializable + json_string = json.dumps(chat_message_data) + assert isinstance(json_string, str) + + # Test round-trip serialization (dump -> load) + restored_chat_message = TextMessage.load(chat_message_data) + assert restored_chat_message.source == "test_agent" + assert restored_chat_message.content == "Hello, world!" + assert restored_chat_message.created_at == test_datetime + + # Test BaseAgentEvent subclass with datetime + agent_event = ModelClientStreamingChunkEvent(source="test_agent", content="chunk", created_at=test_datetime) + + # Test that dump() returns JSON-serializable data + agent_event_data = agent_event.dump() + + # Verify that the datetime is converted to a string in ISO format + assert isinstance(agent_event_data["created_at"], str) + assert agent_event_data["created_at"] == expected_iso + + # Verify that the dumped data is JSON serializable + json_string = json.dumps(agent_event_data) + assert isinstance(json_string, str) + + # Test round-trip serialization (dump -> load) + restored_agent_event = ModelClientStreamingChunkEvent.load(agent_event_data) + assert restored_agent_event.source == "test_agent" + assert restored_agent_event.content == "chunk" + assert restored_agent_event.created_at == test_datetime + + # Test with auto-generated datetime (default created_at) + auto_message = TextMessage(source="test_agent", content="Auto datetime test") + auto_message_data = auto_message.dump() + + # Verify datetime is serialized as string + assert isinstance(auto_message_data["created_at"], str) + + # Verify JSON serialization works without errors + json.dumps(auto_message_data) + + # Test round-trip with auto-generated datetime + restored_auto_message = TextMessage.load(auto_message_data) + assert restored_auto_message.created_at == auto_message.created_at diff --git a/python/packages/autogen-agentchat/tests/test_society_of_mind_agent.py b/python/packages/autogen-agentchat/tests/test_society_of_mind_agent.py index 9634c3d18f71..875fa06a9335 100644 --- a/python/packages/autogen-agentchat/tests/test_society_of_mind_agent.py +++ b/python/packages/autogen-agentchat/tests/test_society_of_mind_agent.py @@ -1,11 +1,15 @@ -from typing import AsyncGenerator +from types import MethodType +from typing import Any, AsyncGenerator, List, Sequence import pytest import pytest_asyncio from autogen_agentchat.agents import AssistantAgent, SocietyOfMindAgent -from autogen_agentchat.conditions import MaxMessageTermination +from autogen_agentchat.base import TaskResult +from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination +from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage, TextMessage from autogen_agentchat.teams import RoundRobinGroupChat from autogen_core import AgentRuntime, SingleThreadedAgentRuntime +from autogen_core.models import CreateResult, LLMMessage, SystemMessage from autogen_ext.models.replay import ReplayChatCompletionClient @@ -31,11 +35,9 @@ async def test_society_of_mind_agent(runtime: AgentRuntime | None) -> None: inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination, runtime=runtime) society_of_mind_agent = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=model_client) response = await society_of_mind_agent.run(task="Count to 10.") - assert len(response.messages) == 4 + assert len(response.messages) == 2 assert response.messages[0].source == "user" - assert response.messages[1].source == "assistant1" - assert response.messages[2].source == "assistant2" - assert response.messages[3].source == "society_of_mind" + assert response.messages[1].source == "society_of_mind" # Test save and load state. state = await society_of_mind_agent.save_state() @@ -57,3 +59,263 @@ async def test_society_of_mind_agent(runtime: AgentRuntime | None) -> None: loaded_soc_agent = SocietyOfMindAgent.load_component(soc_agent_config) assert isinstance(loaded_soc_agent, SocietyOfMindAgent) assert loaded_soc_agent.name == "society_of_mind" + + +@pytest.mark.asyncio +async def test_society_of_mind_agent_output_task_messages_parameter(runtime: AgentRuntime | None) -> None: + """Test that output_task_messages parameter controls whether task messages are included in the stream.""" + model_client = ReplayChatCompletionClient( + ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"], + ) + agent1 = AssistantAgent("assistant1", model_client=model_client, system_message="You are a helpful assistant.") + agent2 = AssistantAgent("assistant2", model_client=model_client, system_message="You are a helpful assistant.") + inner_termination = MaxMessageTermination(2) # Reduce to 2 to use fewer responses + inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination, runtime=runtime) + + # Test 1: Test team with output_task_messages=True (default behavior) + messages_with_task: List[BaseAgentEvent | BaseChatMessage] = [] + async for message in inner_team.run_stream(task="Count to 10", output_task_messages=True): + if not isinstance(message, TaskResult): + messages_with_task.append(message) + + # Should include the task message + assert len(messages_with_task) >= 1 + assert any( + isinstance(msg, TextMessage) and msg.source == "user" and "Count to 10" in msg.content + for msg in messages_with_task + ) + + # Reset team before next test + await inner_team.reset() + + # Test 2: Test team with output_task_messages=False + messages_without_task: List[BaseAgentEvent | BaseChatMessage] = [] + async for message in inner_team.run_stream(task="Count to 10", output_task_messages=False): + if not isinstance(message, TaskResult): + messages_without_task.append(message) + + # Should NOT include the task message in the stream + assert not any( + isinstance(msg, TextMessage) and msg.source == "user" and "Count to 10" in msg.content + for msg in messages_without_task + ) + + # Reset team before next test + await inner_team.reset() + + # Test 3: Test SocietyOfMindAgent uses output_task_messages=False internally + # Create a separate model client for SocietyOfMindAgent to ensure we have enough responses + soma_model_client = ReplayChatCompletionClient( + ["Final response from society of mind"], + ) + society_of_mind_agent = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=soma_model_client) + + # Collect all messages from the SocietyOfMindAgent stream + soma_messages: List[BaseAgentEvent | BaseChatMessage] = [] + async for message in society_of_mind_agent.run_stream(task="Count to 10"): + if not isinstance(message, TaskResult): + soma_messages.append(message) + + # The SocietyOfMindAgent should output the task message (since it's the outer agent) + # but should NOT forward the task messages from its inner team + task_messages_in_soma = [msg for msg in soma_messages if isinstance(msg, TextMessage) and msg.source == "user"] + + # Count how many times "Count to 10" appears in the stream + # With proper implementation, it should appear exactly once (from outer level only) + count_task_messages = sum( + 1 + for msg in soma_messages + if isinstance(msg, TextMessage) and msg.source == "user" and "Count to 10" in msg.content + ) + + # Should have exactly one task message (from the outer level only) + assert len(task_messages_in_soma) == 1 + assert count_task_messages == 1 # Should appear exactly once, not duplicated from inner team + + # Should have the SocietyOfMindAgent's final response + soma_responses = [msg for msg in soma_messages if isinstance(msg, TextMessage) and msg.source == "society_of_mind"] + assert len(soma_responses) == 1 + + +@pytest.mark.asyncio +async def test_society_of_mind_agent_empty_messges(runtime: AgentRuntime | None) -> None: + model_client = ReplayChatCompletionClient( + ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"], + ) + agent1 = AssistantAgent("assistant1", model_client=model_client, system_message="You are a helpful assistant.") + agent2 = AssistantAgent("assistant2", model_client=model_client, system_message="You are a helpful assistant.") + inner_termination = MaxMessageTermination(3) + inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination, runtime=runtime) + society_of_mind_agent = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=model_client) + response = await society_of_mind_agent.run() + assert len(response.messages) == 1 + assert response.messages[0].source == "society_of_mind" + + +@pytest.mark.asyncio +async def test_society_of_mind_agent_no_response(runtime: AgentRuntime | None) -> None: + model_client = ReplayChatCompletionClient( + ["1", "2", "3"], + ) + agent1 = AssistantAgent("assistant1", model_client=model_client, system_message="You are a helpful assistant.") + agent2 = AssistantAgent("assistant2", model_client=model_client, system_message="You are a helpful assistant.") + inner_termination = MaxMessageTermination(1) # Set to 1 to force no response. + inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination, runtime=runtime) + society_of_mind_agent = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=model_client) + response = await society_of_mind_agent.run(task="Count to 10.") + assert len(response.messages) == 2 + assert response.messages[0].source == "user" + assert response.messages[1].source == "society_of_mind" + assert response.messages[1].to_text() == "No response." + + +@pytest.mark.asyncio +async def test_society_of_mind_agent_multiple_rounds(runtime: AgentRuntime | None) -> None: + model_client = ReplayChatCompletionClient( + ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"], + ) + agent1 = AssistantAgent("assistant1", model_client=model_client, system_message="You are a helpful assistant.") + agent2 = AssistantAgent("assistant2", model_client=model_client, system_message="You are a helpful assistant.") + inner_termination = MaxMessageTermination(3) + inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination, runtime=runtime) + society_of_mind_agent = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=model_client) + response = await society_of_mind_agent.run(task="Count to 10.") + assert len(response.messages) == 2 + assert response.messages[0].source == "user" + assert response.messages[1].source == "society_of_mind" + + # Continue. + response = await society_of_mind_agent.run() + assert len(response.messages) == 1 + assert response.messages[0].source == "society_of_mind" + + # Continue. + response = await society_of_mind_agent.run() + assert len(response.messages) == 1 + assert response.messages[0].source == "society_of_mind" + + +@pytest.mark.asyncio +async def test_society_of_mind_agent_no_multiple_system_messages( + monkeypatch: pytest.MonkeyPatch, runtime: AgentRuntime | None +) -> None: + model_client = ReplayChatCompletionClient(["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"]) + + model_client_soma = ReplayChatCompletionClient( + ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"], + model_info={ + "vision": False, + "function_calling": False, + "json_output": False, + "family": "unknown", + "structured_output": False, + "multiple_system_messages": False, + }, + ) + + original_create = model_client_soma.create + + # mock method with bound self + async def _mock_create( + self: ReplayChatCompletionClient, messages: Sequence[LLMMessage], *args: Any, **kwargs: Any + ) -> CreateResult: + for message in messages: + assert not isinstance(message, SystemMessage) + kwargs["messages"] = messages + return await original_create(*args, **kwargs) + + # bind it + monkeypatch.setattr(model_client_soma, "create", MethodType(_mock_create, model_client_soma)) + + agent1 = AssistantAgent("assistant1", model_client=model_client, system_message="You are a helpful assistant.") + agent2 = AssistantAgent("assistant2", model_client=model_client, system_message="You are a helpful assistant.") + inner_termination = MaxMessageTermination(3) + inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination, runtime=runtime) + society_of_mind_agent = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=model_client_soma) + await society_of_mind_agent.run(task="Count to 10.") + + +@pytest.mark.asyncio +async def test_society_of_mind_agent_yes_multiple_system_messages( + monkeypatch: pytest.MonkeyPatch, runtime: AgentRuntime | None +) -> None: + model_client = ReplayChatCompletionClient(["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"]) + + model_client_soma = ReplayChatCompletionClient( + ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"], + model_info={ + "vision": False, + "function_calling": False, + "json_output": False, + "family": "unknown", + "structured_output": False, + "multiple_system_messages": True, + }, + ) + + original_create = model_client_soma.create + + # mock method with bound self + async def _mock_create( + self: ReplayChatCompletionClient, messages: Sequence[LLMMessage], *args: Any, **kwargs: Any + ) -> CreateResult: + assert isinstance(messages[0], SystemMessage) + assert isinstance(messages[-1], SystemMessage) + kwargs["messages"] = messages + return await original_create(*args, **kwargs) + + # bind it + monkeypatch.setattr(model_client_soma, "create", MethodType(_mock_create, model_client_soma)) + + agent1 = AssistantAgent("assistant1", model_client=model_client, system_message="You are a helpful assistant.") + agent2 = AssistantAgent("assistant2", model_client=model_client, system_message="You are a helpful assistant.") + inner_termination = MaxMessageTermination(3) + inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination, runtime=runtime) + society_of_mind_agent = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=model_client_soma) + await society_of_mind_agent.run(task="Count to 10.") + + +@pytest.mark.asyncio +async def test_default_output_task_messages_behavior() -> None: + """Test that task messages are included by default (backward compatibility).""" + # Create inner team + model_client = ReplayChatCompletionClient(["Hello", "World", "TERMINATE"]) + agent1 = AssistantAgent("agent1", model_client=model_client) + agent2 = AssistantAgent("agent2", model_client=model_client) + termination = TextMentionTermination("TERMINATE") + inner_team = RoundRobinGroupChat(participants=[agent1, agent2], termination_condition=termination) + + streamed_messages: List[BaseAgentEvent | BaseChatMessage] = [] + final_result: TaskResult | None = None + + # Test default behavior (should include task messages since default is True) + async for message in inner_team.run_stream(task="Test default behavior"): + if isinstance(message, TaskResult): + final_result = message + else: + streamed_messages.append(message) + + # Verify default behavior: task message should be included in stream + assert final_result is not None + task_message_found_in_stream = any( + isinstance(msg, TextMessage) and msg.source == "user" and "Test default behavior" in msg.content + for msg in streamed_messages + ) + assert task_message_found_in_stream, "Task message should be included in stream by default" + + # Validate that task message is included in the TaskResult.messages by default + task_message_in_result = any( + isinstance(msg, TextMessage) and msg.source == "user" and "Test default behavior" in msg.content + for msg in final_result.messages + ) + assert task_message_in_result, "Task message should be included in TaskResult.messages by default" + + # Verify the content structure makes sense (task message + agent responses) + user_messages = [msg for msg in final_result.messages if isinstance(msg, TextMessage) and msg.source == "user"] + agent_messages = [ + msg for msg in final_result.messages if isinstance(msg, TextMessage) and msg.source in ["agent1", "agent2"] + ] + + assert len(user_messages) >= 1, "Should have at least one user message (the task)" + assert len(agent_messages) >= 1, "Should have at least one agent response" + assert user_messages[0].content == "Test default behavior", "First user message should be the task" diff --git a/python/packages/autogen-agentchat/tests/test_streaming_message_id_correlation.py b/python/packages/autogen-agentchat/tests/test_streaming_message_id_correlation.py new file mode 100644 index 000000000000..9031c532a6ea --- /dev/null +++ b/python/packages/autogen-agentchat/tests/test_streaming_message_id_correlation.py @@ -0,0 +1,99 @@ +from typing import List, Optional + +import pytest +from autogen_agentchat.agents import AssistantAgent +from autogen_agentchat.base import TaskResult +from autogen_agentchat.messages import ModelClientStreamingChunkEvent, TextMessage +from autogen_core import FunctionCall +from autogen_core.models import CreateResult, ModelFamily, RequestUsage +from autogen_ext.models.replay import ReplayChatCompletionClient + + +async def _echo_function(input: str) -> str: + return input + + +@pytest.mark.asyncio +async def test_streaming_message_id_correlation() -> None: + """Test that streaming chunks have full_message_id that matches final message ID.""" + mock_client = ReplayChatCompletionClient( + [ + "Response to message", + ] + ) + agent = AssistantAgent( + "test_agent", + model_client=mock_client, + model_client_stream=True, + ) + + # Track all chunks and the final message + chunks: List[ModelClientStreamingChunkEvent] = [] + final_message: Optional[TextMessage] = None + + async for message in agent.run_stream(task="task"): + if isinstance(message, TaskResult): + assert len(message.messages) == 2 + assert isinstance(message.messages[0], TextMessage) + assert isinstance(message.messages[1], TextMessage) + final_message = message.messages[1] + elif isinstance(message, ModelClientStreamingChunkEvent): + chunks.append(message) + + # Verify we got chunks and a final message + assert len(chunks) > 0 + assert final_message is not None + + # Every chunk should have the same full_message_id as the final message's id + for chunk in chunks: + assert chunk.full_message_id == final_message.id + + # Test the reflect_on_tool_use streaming case + mock_client = ReplayChatCompletionClient( + [ + CreateResult( + content=[ + FunctionCall(id="1", name="_echo_function", arguments=r'{"input": "task"}'), + ], + finish_reason="function_calls", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ), + "Example reflection response", + ], + model_info={ + "function_calling": True, + "vision": False, + "json_output": False, + "family": ModelFamily.GPT_4, + "structured_output": False, + }, + ) + + agent = AssistantAgent( + "test_agent", + model_client=mock_client, + model_client_stream=True, + reflect_on_tool_use=True, + tools=[_echo_function], + ) + + # Track reflection chunks and final message + reflection_chunks: List[ModelClientStreamingChunkEvent] = [] + final_reflection_message: Optional[TextMessage] = None + + async for message in agent.run_stream(task="task"): + if isinstance(message, TaskResult): + # The last message should be the reflection result + if isinstance(message.messages[-1], TextMessage): + final_reflection_message = message.messages[-1] + elif isinstance(message, ModelClientStreamingChunkEvent): + reflection_chunks.append(message) + + # Verify we got reflection chunks and a final message + assert len(reflection_chunks) > 0 + assert final_reflection_message is not None + + # Every reflection chunk should have the same full_message_id as the final message's id + for chunk in reflection_chunks: + assert chunk.full_message_id == final_reflection_message.id # type: ignore diff --git a/python/packages/autogen-agentchat/tests/test_task_runner_tool.py b/python/packages/autogen-agentchat/tests/test_task_runner_tool.py new file mode 100644 index 000000000000..2449524593eb --- /dev/null +++ b/python/packages/autogen-agentchat/tests/test_task_runner_tool.py @@ -0,0 +1,239 @@ +import pytest +from autogen_agentchat.agents import AssistantAgent +from autogen_agentchat.conditions import MaxMessageTermination +from autogen_agentchat.messages import TextMessage, ToolCallExecutionEvent, ToolCallRequestEvent +from autogen_agentchat.teams import RoundRobinGroupChat +from autogen_agentchat.tools import AgentTool, TeamTool +from autogen_core import ( + CancellationToken, + FunctionCall, +) +from autogen_core.models import CreateResult, RequestUsage +from autogen_ext.models.replay import ReplayChatCompletionClient +from test_group_chat import _EchoAgent # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_agent_tool_run() -> None: + """Test running a task with AgentTool.""" + mock_chat_agent = _EchoAgent("Mock_Agent", "A mock agent for testing") + tool = AgentTool(agent=mock_chat_agent) + task_result = await tool.run_json({"task": "Test task"}, cancellation_token=CancellationToken()) + assert task_result.messages[1].content == "Test task" + + +@pytest.mark.asyncio +async def test_agent_tool_state() -> None: + """Test saving state of AgentTool.""" + mock_chat_agent = _EchoAgent("Mock_Agent", "A mock agent for testing") + tool = AgentTool(agent=mock_chat_agent) + state = await tool.save_state_json() + assert state == {"last_message": None, "total_messages": 0} + + await tool.run_json({"task": "Test task"}, cancellation_token=CancellationToken()) + state = await tool.save_state_json() + assert state == {"last_message": "Test task", "total_messages": 1} + + mock_chat_agent_2 = _EchoAgent("Mock_Agent_2", "A mock agent for testing") + tool_2 = AgentTool(agent=mock_chat_agent_2) + await tool_2.load_state_json(state) + state2 = await tool_2.save_state_json() + assert state2 == {"last_message": "Test task", "total_messages": 1} + + +def test_agent_tool_component() -> None: + """Test serialization of AgentTool to config.""" + model_client = ReplayChatCompletionClient(["test"]) + agent = AssistantAgent(name="assistant", model_client=model_client) + tool = AgentTool(agent=agent) + config = tool.dump_component() + assert config.provider == "autogen_agentchat.tools.AgentTool" + + tool2 = AgentTool.load_component(config) + assert isinstance(tool2, AgentTool) + assert tool2.name == agent.name + assert tool2.description == agent.description + + +@pytest.mark.asyncio +async def test_team_tool() -> None: + """Test running a task with TeamTool.""" + agent1 = _EchoAgent("Agent1", "An agent for testing") + agent2 = _EchoAgent("Agent2", "Another agent for testing") + termination = MaxMessageTermination(max_messages=3) + team = RoundRobinGroupChat( + [agent1, agent2], + termination_condition=termination, + ) + tool = TeamTool(team=team, name="Team Tool", description="A team tool for testing") + task_result = await tool.run_json(args={"task": "test task"}, cancellation_token=CancellationToken()) + assert task_result.messages[1].content == "test task" + assert task_result.messages[2].content == "test task" + + # Validate state. + state = await tool.save_state_json() + # Reload the state and check if it matches. + agent2 = _EchoAgent("Agent1", "Another agent for testing") + agent3 = _EchoAgent("Agent2", "Another agent for testing") + team2 = RoundRobinGroupChat( + [agent2, agent3], + termination_condition=termination, + ) + tool2 = TeamTool(team=team2, name="Team Tool", description="A team tool for testing") + await tool2.load_state_json(state) + state2 = await tool2.save_state_json() + assert state == state2 + + +@pytest.mark.asyncio +async def test_team_tool_component() -> None: + """Test serialization of TeamTool to config.""" + model_client = ReplayChatCompletionClient(["test"]) + agent1 = AssistantAgent(name="assistant1", model_client=model_client) + agent2 = AssistantAgent(name="assistant2", model_client=model_client) + team = RoundRobinGroupChat([agent1, agent2]) + tool = TeamTool(team=team, name="Team Tool", description="A team tool for testing") + config = tool.dump_component() + assert config.provider == "autogen_agentchat.tools.TeamTool" + + tool2 = TeamTool.load_component(config) + assert isinstance(tool2, TeamTool) + assert tool2.name == "Team Tool" + assert tool2.description == "A team tool for testing" + assert isinstance(tool2._team, RoundRobinGroupChat) # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_agent_tool_stream() -> None: + """Test running a task with AgentTool in streaming mode.""" + + def _query_function() -> str: + return "Test task" + + tool_agent_model_client = ReplayChatCompletionClient( + [ + CreateResult( + content=[FunctionCall(name="query_function", arguments="{}", id="1")], + finish_reason="function_calls", + usage=RequestUsage(prompt_tokens=0, completion_tokens=0), + cached=False, + ), + "Summary from tool agent", + ], + model_info={ + "family": "gpt-41", + "function_calling": True, + "json_output": True, + "multiple_system_messages": True, + "structured_output": True, + "vision": True, + }, + ) + tool_agent = AssistantAgent( + name="tool_agent", + model_client=tool_agent_model_client, + tools=[_query_function], + reflect_on_tool_use=True, + description="An agent for testing", + ) + tool = AgentTool(tool_agent) + + main_agent_model_client = ReplayChatCompletionClient( + [ + CreateResult( + content=[FunctionCall(id="1", name="tool_agent", arguments='{"task": "Input task from main agent"}')], + finish_reason="function_calls", + usage=RequestUsage(prompt_tokens=0, completion_tokens=0), + cached=False, + ), + "Summary from main agent", + ], + model_info={ + "family": "gpt-41", + "function_calling": True, + "json_output": True, + "multiple_system_messages": True, + "structured_output": True, + "vision": True, + }, + ) + + main_agent = AssistantAgent( + name="main_agent", + model_client=main_agent_model_client, + tools=[tool], + reflect_on_tool_use=True, + description="An agent for testing", + ) + result = await main_agent.run(task="Input task from user", cancellation_token=CancellationToken()) + assert isinstance(result.messages[0], TextMessage) + assert result.messages[0].content == "Input task from user" + assert isinstance(result.messages[1], ToolCallRequestEvent) + assert isinstance(result.messages[2], TextMessage) + assert result.messages[2].content == "Input task from main agent" + assert isinstance(result.messages[3], ToolCallRequestEvent) + assert isinstance(result.messages[4], ToolCallExecutionEvent) + assert isinstance(result.messages[5], TextMessage) + assert result.messages[5].content == "Summary from tool agent" + assert isinstance(result.messages[6], ToolCallExecutionEvent) + assert result.messages[6].content[0].content == "tool_agent: Summary from tool agent" + assert isinstance(result.messages[7], TextMessage) + assert result.messages[7].content == "Summary from main agent" + + +@pytest.mark.asyncio +async def test_team_tool_stream() -> None: + """Test running a task with TeamTool in streaming mode.""" + agent1 = _EchoAgent("Agent1", "An agent for testing") + agent2 = _EchoAgent("Agent2", "Another agent for testing") + termination = MaxMessageTermination(max_messages=3) + team = RoundRobinGroupChat( + [agent1, agent2], + termination_condition=termination, + ) + tool = TeamTool( + team=team, name="team_tool", description="A team tool for testing", return_value_as_last_message=True + ) + + model_client = ReplayChatCompletionClient( + [ + CreateResult( + content=[FunctionCall(name="team_tool", arguments='{"task": "test task from main agent"}', id="1")], + finish_reason="function_calls", + usage=RequestUsage(prompt_tokens=0, completion_tokens=0), + cached=False, + ), + "Summary from main agent", + ], + model_info={ + "family": "gpt-41", + "function_calling": True, + "json_output": True, + "multiple_system_messages": True, + "structured_output": True, + "vision": True, + }, + ) + main_agent = AssistantAgent( + name="main_agent", + model_client=model_client, + tools=[tool], + reflect_on_tool_use=True, + description="An agent for testing", + ) + result = await main_agent.run(task="test task from user", cancellation_token=CancellationToken()) + assert isinstance(result.messages[0], TextMessage) + assert result.messages[0].content == "test task from user" + assert isinstance(result.messages[1], ToolCallRequestEvent) + assert isinstance(result.messages[2], TextMessage) + assert result.messages[2].content == "test task from main agent" + assert isinstance(result.messages[3], TextMessage) + assert result.messages[3].content == "test task from main agent" + assert result.messages[3].source == "Agent1" + assert isinstance(result.messages[4], TextMessage) + assert result.messages[4].content == "test task from main agent" + assert result.messages[4].source == "Agent2" + assert isinstance(result.messages[5], ToolCallExecutionEvent) + assert result.messages[5].content[0].content == "test task from main agent" + assert isinstance(result.messages[6], TextMessage) + assert result.messages[6].content == "Summary from main agent" diff --git a/python/packages/autogen-agentchat/tests/test_termination_condition.py b/python/packages/autogen-agentchat/tests/test_termination_condition.py index 230a8d4ac721..bccb2f012499 100644 --- a/python/packages/autogen-agentchat/tests/test_termination_condition.py +++ b/python/packages/autogen-agentchat/tests/test_termination_condition.py @@ -1,9 +1,11 @@ import asyncio +from typing import Sequence import pytest from autogen_agentchat.base import TerminatedException from autogen_agentchat.conditions import ( ExternalTermination, + FunctionalTermination, FunctionCallTermination, HandoffTermination, MaxMessageTermination, @@ -15,13 +17,17 @@ TokenUsageTermination, ) from autogen_agentchat.messages import ( + BaseAgentEvent, + BaseChatMessage, HandoffMessage, StopMessage, + StructuredMessage, TextMessage, ToolCallExecutionEvent, UserInputRequestedEvent, ) from autogen_core.models import FunctionExecutionResult, RequestUsage +from pydantic import BaseModel @pytest.mark.asyncio @@ -375,3 +381,53 @@ async def test_function_call_termination() -> None: ) assert not termination.terminated await termination.reset() + + +@pytest.mark.asyncio +async def test_functional_termination() -> None: + async def async_termination_func(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> bool: + if len(messages) < 1: + return False + if isinstance(messages[-1], TextMessage): + return messages[-1].content == "stop" + return False + + termination = FunctionalTermination(async_termination_func) + assert await termination([]) is None + await termination.reset() + + assert await termination([TextMessage(content="Hello", source="user")]) is None + await termination.reset() + + assert await termination([TextMessage(content="stop", source="user")]) is not None + assert termination.terminated + await termination.reset() + + assert await termination([TextMessage(content="Hello", source="user")]) is None + + class TestContentType(BaseModel): + content: str + data: str + + def sync_termination_func(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> bool: + if len(messages) < 1: + return False + last_message = messages[-1] + if isinstance(last_message, StructuredMessage) and isinstance(last_message.content, TestContentType): # type: ignore[reportUnknownMemberType] + return last_message.content.data == "stop" + return False + + termination = FunctionalTermination(sync_termination_func) + assert await termination([]) is None + await termination.reset() + assert await termination([TextMessage(content="Hello", source="user")]) is None + await termination.reset() + assert ( + await termination( + [StructuredMessage[TestContentType](content=TestContentType(content="1", data="stop"), source="user")] + ) + is not None + ) + assert termination.terminated + await termination.reset() + assert await termination([TextMessage(content="Hello", source="user")]) is None diff --git a/python/packages/autogen-agentchat/tests/test_userproxy_agent.py b/python/packages/autogen-agentchat/tests/test_userproxy_agent.py index 8ff6422a696f..855211de82a9 100644 --- a/python/packages/autogen-agentchat/tests/test_userproxy_agent.py +++ b/python/packages/autogen-agentchat/tests/test_userproxy_agent.py @@ -4,7 +4,7 @@ import pytest from autogen_agentchat.agents import UserProxyAgent from autogen_agentchat.base import Response -from autogen_agentchat.messages import ChatMessage, HandoffMessage, TextMessage +from autogen_agentchat.messages import BaseChatMessage, HandoffMessage, TextMessage from autogen_core import CancellationToken @@ -53,7 +53,7 @@ def custom_input(prompt: str) -> str: agent = UserProxyAgent(name="test_user", input_func=custom_input) - messages: Sequence[ChatMessage] = [ + messages: Sequence[BaseChatMessage] = [ TextMessage(content="Initial message", source="assistant"), HandoffMessage(content="Handing off to user for confirmation", source="assistant", target="test_user"), ] diff --git a/python/packages/autogen-agentchat/tests/utils.py b/python/packages/autogen-agentchat/tests/utils.py index 90d66f0ba983..85d0d41f189d 100644 --- a/python/packages/autogen-agentchat/tests/utils.py +++ b/python/packages/autogen-agentchat/tests/utils.py @@ -2,7 +2,10 @@ import logging import sys from datetime import datetime +from typing import Sequence +from autogen_agentchat.base._task import TaskResult +from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage, BaseTextChatMessage from pydantic import BaseModel @@ -18,7 +21,7 @@ def emit(self, record: logging.LogRecord) -> None: record.msg = json.dumps( { "timestamp": ts, - "message": record.msg.model_dump(), + "message": record.msg.model_dump_json(indent=2), "type": record.msg.__class__.__name__, }, ) @@ -37,3 +40,36 @@ def emit(self, record: logging.LogRecord) -> None: }, ) sys.stdout.write(f"{record.msg}\n") + + +def compare_messages( + msg1: BaseAgentEvent | BaseChatMessage | BaseTextChatMessage, + msg2: BaseAgentEvent | BaseChatMessage | BaseTextChatMessage, +) -> bool: + if isinstance(msg1, BaseTextChatMessage) and isinstance(msg2, BaseTextChatMessage): + if msg1.content != msg2.content: + return False + return ( + (msg1.source == msg2.source) and (msg1.models_usage == msg2.models_usage) and (msg1.metadata == msg2.metadata) + ) + + +def compare_message_lists( + msgs1: Sequence[BaseAgentEvent | BaseChatMessage], + msgs2: Sequence[BaseAgentEvent | BaseChatMessage], +) -> bool: + if len(msgs1) != len(msgs2): + return False + for i in range(len(msgs1)): + if not compare_messages(msgs1[i], msgs2[i]): + return False + return True + + +def compare_task_results( + res1: TaskResult, + res2: TaskResult, +) -> bool: + if res1.stop_reason != res2.stop_reason: + return False + return compare_message_lists(res1.messages, res2.messages) diff --git a/python/packages/autogen-core/.gitignore b/python/packages/autogen-core/.gitignore index 06f543fa9140..a93071a96970 100644 --- a/python/packages/autogen-core/.gitignore +++ b/python/packages/autogen-core/.gitignore @@ -170,4 +170,4 @@ log.jsonl docs/**/jupyter_execute # Temporary files -tmp_code_*.py +tmp_code_*.py \ No newline at end of file diff --git a/python/packages/autogen-core/docs/README.md b/python/packages/autogen-core/docs/README.md deleted file mode 100644 index 5aff221eaa50..000000000000 --- a/python/packages/autogen-core/docs/README.md +++ /dev/null @@ -1,37 +0,0 @@ -## Building the AutoGen Documentation - -AutoGen documentation is based on the sphinx documentation system and uses the myst-parser to render markdown files. It uses the [pydata-sphinx-theme](https://pydata-sphinx-theme.readthedocs.io/en/latest/) to style the documentation. - -### Prerequisites - -Ensure you have all of the dev dependencies for the `autogen-core` package installed. You can install them by running the following command from the root of the python repository: - -```bash -uv sync -source .venv/bin/activate -``` - -## Building Docs - -To build the documentation, run the following command from the root of the python repository: - -```bash -poe --directory ./packages/autogen-core/ docs-build -``` - -To serve the documentation locally, run the following command from the root of the python repository: - -```bash -poe --directory ./packages/autogen-core/ docs-serve -``` - -[!NOTE] -Sphinx will only rebuild files that have changed since the last build. If you want to force a full rebuild, you can delete the `./packages/autogen-core/docs/build` directory before running the `docs-build` command. - -## Versioning the Documentation - -The current theme - [pydata-sphinx-theme](https://pydata-sphinx-theme.readthedocs.io/en/latest/) - supports [switching between versions](https://pydata-sphinx-theme.readthedocs.io/en/stable/user_guide/version-dropdown.html) of the documentation. - -To version the documentation, you need to create a new version of the documentation by copying the existing documentation to a new directory with the version number. For example, to create a new version of the documentation for version `0.1.0`, you would run the following command: - -How are various versions built? - TBD. diff --git a/python/packages/autogen-core/docs/src/images/assistant-agent.svg b/python/packages/autogen-core/docs/src/images/assistant-agent.svg deleted file mode 100644 index fc1ce4ba1b38..000000000000 --- a/python/packages/autogen-core/docs/src/images/assistant-agent.svg +++ /dev/null @@ -1,3 +0,0 @@ - - -
Model Context
Model Context
Model Context
Model Context
New Messages
New Messages
Memory
Memory
Model Client
Model Client
Tools
Tools
Model Context
Model Context
1. Add New Messages to Context
1. Add New Messages to Context
Model Context
Model Context
2. Update Context
2. Update Context
Model Context
Model Context
3. Chat Completion
3. Chat Completion
Model Context
Model Context
Model Context
Model Context
Model Context
Model Context
4. Tool Execution
4. Tool Execution
Model Client
Model Client
5. Chat Completion (Reflect on Tool Use)
5. Chat Completion (R...
Model Context
Model Context
Model Context
Model Context
Model Context
Model Context
Model Context
Model Context
Response
(Text)
Response...
Response 
(Tool Result Summary)
Response...
Response 
(Text Message)
Response...
Model Context
Model Context
Model Context
Model Context
Model Context
Model Context
Tool Call Detected?
Tool Call Detected?
Reflect on Tool Use?
Reflect on Tool Use?
No
No
No
No
Handoff Detected?
Handoff Detected?
Response
(Handoff)
Response...
Yes
Yes
Assistant Agent
Assistant Agent \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/reference/index.md b/python/packages/autogen-core/docs/src/reference/index.md deleted file mode 100644 index d498a60fc9c5..000000000000 --- a/python/packages/autogen-core/docs/src/reference/index.md +++ /dev/null @@ -1,74 +0,0 @@ ---- -myst: - html_meta: - "description lang=en": | - AutoGen is a community-driven project. Learn how to get involved, contribute, and connect with the community. ---- - -# API Reference - -```{toctree} -:caption: AutoGen AgentChat -:maxdepth: 2 - -python/autogen_agentchat -python/autogen_agentchat.messages -python/autogen_agentchat.agents -python/autogen_agentchat.teams -python/autogen_agentchat.base -python/autogen_agentchat.conditions -python/autogen_agentchat.ui -python/autogen_agentchat.state -``` - -```{toctree} -:caption: AutoGen Core -:maxdepth: 2 - -python/autogen_core -python/autogen_core.code_executor -python/autogen_core.models -python/autogen_core.model_context -python/autogen_core.tools -python/autogen_core.tool_agent -python/autogen_core.memory -python/autogen_core.exceptions -python/autogen_core.logging -``` - -```{toctree} -:caption: AutoGen Extensions -:maxdepth: 2 - -python/autogen_ext.agents.magentic_one -python/autogen_ext.agents.openai -python/autogen_ext.agents.web_surfer -python/autogen_ext.agents.file_surfer -python/autogen_ext.agents.video_surfer -python/autogen_ext.agents.video_surfer.tools -python/autogen_ext.teams.magentic_one -python/autogen_ext.models.cache -python/autogen_ext.models.openai -python/autogen_ext.models.replay -python/autogen_ext.models.azure -python/autogen_ext.models.anthropic -python/autogen_ext.models.semantic_kernel -python/autogen_ext.models.ollama -python/autogen_ext.models.llama_cpp -python/autogen_ext.tools.code_execution -python/autogen_ext.tools.graphrag -python/autogen_ext.tools.http -python/autogen_ext.tools.langchain -python/autogen_ext.tools.mcp -python/autogen_ext.tools.semantic_kernel -python/autogen_ext.code_executors.local -python/autogen_ext.code_executors.docker -python/autogen_ext.code_executors.jupyter -python/autogen_ext.code_executors.azure -python/autogen_ext.cache_store.diskcache -python/autogen_ext.cache_store.redis -python/autogen_ext.runtimes.grpc -python/autogen_ext.auth.azure -python/autogen_ext.experimental.task_centric_memory -python/autogen_ext.experimental.task_centric_memory.utils -``` diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.agents.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.agents.rst deleted file mode 100644 index 4ccb296c4aec..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.agents.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_agentchat.agents -========================= - - -.. automodule:: autogen_agentchat.agents - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.base.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.base.rst deleted file mode 100644 index c973f560ac6b..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.base.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_agentchat.base -======================= - - -.. automodule:: autogen_agentchat.base - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.conditions.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.conditions.rst deleted file mode 100644 index 09db1eeb4773..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.conditions.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_agentchat.conditions -============================= - - -.. automodule:: autogen_agentchat.conditions - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.messages.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.messages.rst deleted file mode 100644 index b1dcf2932197..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.messages.rst +++ /dev/null @@ -1,7 +0,0 @@ -autogen\_agentchat.messages -=========================== - -.. automodule:: autogen_agentchat.messages - :members: - :undoc-members: - :show-inheritance: \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.rst deleted file mode 100644 index bf5b697366a1..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.rst +++ /dev/null @@ -1,7 +0,0 @@ -autogen\_agentchat -================== - -.. automodule:: autogen_agentchat - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.state.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.state.rst deleted file mode 100644 index 72127b8bab80..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.state.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_agentchat.state -======================== - - -.. automodule:: autogen_agentchat.state - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.teams.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.teams.rst deleted file mode 100644 index cfa76107ee08..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.teams.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_agentchat.teams -======================== - - -.. automodule:: autogen_agentchat.teams - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.ui.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.ui.rst deleted file mode 100644 index 7d5db1029788..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_agentchat.ui.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_agentchat.ui -===================== - - -.. automodule:: autogen_agentchat.ui - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_core.code_executor.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_core.code_executor.rst deleted file mode 100644 index f05bdd68d9f5..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_core.code_executor.rst +++ /dev/null @@ -1,7 +0,0 @@ -autogen\_core.code_executor -=========================== - -.. automodule:: autogen_core.code_executor - :members: - :undoc-members: - :show-inheritance: \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_core.exceptions.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_core.exceptions.rst deleted file mode 100644 index 5ca415524b16..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_core.exceptions.rst +++ /dev/null @@ -1,7 +0,0 @@ -autogen\_core.exceptions -======================== - -.. automodule:: autogen_core.exceptions - :members: - :undoc-members: - :show-inheritance: \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_core.logging.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_core.logging.rst deleted file mode 100644 index 6caaf7475e30..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_core.logging.rst +++ /dev/null @@ -1,7 +0,0 @@ -autogen\_core.logging -===================== - -.. automodule:: autogen_core.logging - :members: - :undoc-members: - :show-inheritance: \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_core.memory.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_core.memory.rst deleted file mode 100644 index c2f7f446cd49..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_core.memory.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_core.memory -==================== - - -.. automodule:: autogen_core.memory - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_core.model_context.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_core.model_context.rst deleted file mode 100644 index e8beee6b2ca0..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_core.model_context.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_core.model\_context -======================================= - - -.. automodule:: autogen_core.model_context - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_core.models.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_core.models.rst deleted file mode 100644 index 9a2b06cfd46e..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_core.models.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_core.models -=============================== - - -.. automodule:: autogen_core.models - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_core.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_core.rst deleted file mode 100644 index c3f7bded900d..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_core.rst +++ /dev/null @@ -1,11 +0,0 @@ -autogen\_core -============= - -.. automodule:: autogen_core - :members: - :undoc-members: - :show-inheritance: - :member-order: bysource - -.. For some reason the following isn't picked up by automodule above -.. autoclass:: autogen_core.ComponentType \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_core.tool_agent.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_core.tool_agent.rst deleted file mode 100644 index 18e8e8a7df5b..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_core.tool_agent.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_core.tool\_agent -==================================== - - -.. automodule:: autogen_core.tool_agent - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_core.tools.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_core.tools.rst deleted file mode 100644 index c78d6a750614..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_core.tools.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_core.tools -=================== - - -.. automodule:: autogen_core.tools - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.file_surfer.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.file_surfer.rst deleted file mode 100644 index fb83308cecf3..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.file_surfer.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.agents.file\_surfer -================================ - - -.. automodule:: autogen_ext.agents.file_surfer - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.magentic_one.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.magentic_one.rst deleted file mode 100644 index 93839ba881da..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.magentic_one.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.agents.magentic\_one -================================= - - -.. automodule:: autogen_ext.agents.magentic_one - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.openai.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.openai.rst deleted file mode 100644 index e3d068d71cbd..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.openai.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.agents.openai -================================ - - -.. automodule:: autogen_ext.agents.openai - :members: - :undoc-members: - :show-inheritance: \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.video_surfer.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.video_surfer.rst deleted file mode 100644 index a171573c2ea3..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.video_surfer.rst +++ /dev/null @@ -1,7 +0,0 @@ -autogen\_ext.agents.video\_surfer -================================= - -.. automodule:: autogen_ext.agents.video_surfer - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.video_surfer.tools.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.video_surfer.tools.rst deleted file mode 100644 index 1fe80764b3e2..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.video_surfer.tools.rst +++ /dev/null @@ -1,7 +0,0 @@ -autogen\_ext.agents.video\_surfer.tools -======================================= - -.. automodule:: autogen_ext.agents.video_surfer.tools - :members: - :undoc-members: - :show-inheritance: \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.web_surfer.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.web_surfer.rst deleted file mode 100644 index 8ae0502ae4dc..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.agents.web_surfer.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.agents.web\_surfer -=============================== - - -.. automodule:: autogen_ext.agents.web_surfer - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.auth.azure.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.auth.azure.rst deleted file mode 100644 index 7d27398260df..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.auth.azure.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.auth.azure -======================= - - -.. automodule:: autogen_ext.auth.azure - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.cache_store.diskcache.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.cache_store.diskcache.rst deleted file mode 100644 index 5fbc4c8b35ac..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.cache_store.diskcache.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.cache_store.diskcache -================================== - - -.. automodule:: autogen_ext.cache_store.diskcache - :members: - :undoc-members: - :show-inheritance: \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.cache_store.redis.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.cache_store.redis.rst deleted file mode 100644 index fab1b46d520a..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.cache_store.redis.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.cache_store.redis -============================== - - -.. automodule:: autogen_ext.cache_store.redis - :members: - :undoc-members: - :show-inheritance: \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.code_executors.azure.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.code_executors.azure.rst deleted file mode 100644 index 08fd3bdcfb7a..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.code_executors.azure.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.code\_executors.azure -================================== - - -.. automodule:: autogen_ext.code_executors.azure - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.code_executors.docker.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.code_executors.docker.rst deleted file mode 100644 index b4a9a287a820..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.code_executors.docker.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.code\_executors.docker -=================================== - - -.. automodule:: autogen_ext.code_executors.docker - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.code_executors.jupyter.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.code_executors.jupyter.rst deleted file mode 100644 index 6b44c137f810..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.code_executors.jupyter.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.code\_executors.jupyter -==================================== - - -.. automodule:: autogen_ext.code_executors.jupyter - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.code_executors.local.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.code_executors.local.rst deleted file mode 100644 index 7fa6f06d6a89..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.code_executors.local.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.code\_executors.local -================================== - - -.. automodule:: autogen_ext.code_executors.local - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.experimental.task_centric_memory.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.experimental.task_centric_memory.rst deleted file mode 100644 index e4bf80c70abb..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.experimental.task_centric_memory.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.experimental.task\_centric\_memory -=============================================== - - -.. automodule:: autogen_ext.experimental.task_centric_memory - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.experimental.task_centric_memory.utils.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.experimental.task_centric_memory.utils.rst deleted file mode 100644 index c5f7a58bfead..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.experimental.task_centric_memory.utils.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.experimental.task\_centric\_memory.utils -===================================================== - - -.. automodule:: autogen_ext.experimental.task_centric_memory.utils - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.anthropic.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.anthropic.rst deleted file mode 100644 index 51d882829eed..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.anthropic.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.models.anthropic -============================= - - -.. automodule:: autogen_ext.models.anthropic - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.azure.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.azure.rst deleted file mode 100644 index 64c16a5a57d4..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.azure.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.models.azure -========================== - - -.. automodule:: autogen_ext.models.azure - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.cache.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.cache.rst deleted file mode 100644 index 48956ace16e2..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.cache.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.models.cache -========================= - - -.. automodule:: autogen_ext.models.cache - :members: - :undoc-members: - :show-inheritance: \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.llama_cpp.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.llama_cpp.rst deleted file mode 100644 index 56470d36d646..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.llama_cpp.rst +++ /dev/null @@ -1,9 +0,0 @@ -autogen\_ext.models.llama\_cpp -============================== - - -.. automodule:: autogen_ext.models.llama_cpp - :members: - :undoc-members: - :show-inheritance: - :member-order: bysource diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.ollama.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.ollama.rst deleted file mode 100644 index 67076e797251..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.ollama.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.models.ollama -========================== - - -.. automodule:: autogen_ext.models.ollama - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.openai.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.openai.rst deleted file mode 100644 index b74e4903f52b..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.openai.rst +++ /dev/null @@ -1,9 +0,0 @@ -autogen\_ext.models.openai -========================== - - -.. automodule:: autogen_ext.models.openai - :members: - :undoc-members: - :show-inheritance: - :member-order: bysource diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.replay.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.replay.rst deleted file mode 100644 index 4fc9aefbfb3d..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.replay.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.models.replay -========================== - - -.. automodule:: autogen_ext.models.replay - :members: - :undoc-members: - :show-inheritance: \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.semantic_kernel.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.semantic_kernel.rst deleted file mode 100644 index c80d462535de..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.models.semantic_kernel.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.models.semantic_kernel -==================================== - - -.. automodule:: autogen_ext.models.semantic_kernel - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.runtimes.grpc.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.runtimes.grpc.rst deleted file mode 100644 index 6a26dd2c3e1f..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.runtimes.grpc.rst +++ /dev/null @@ -1,7 +0,0 @@ -autogen\_ext.runtimes.grpc -========================== - -.. automodule:: autogen_ext.runtimes.grpc - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.teams.magentic_one.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.teams.magentic_one.rst deleted file mode 100644 index 4e18d60a9142..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.teams.magentic_one.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.teams.magentic\_one -================================= - - -.. automodule:: autogen_ext.teams.magentic_one - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.code_execution.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.code_execution.rst deleted file mode 100644 index 18dce6613bfd..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.code_execution.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.tools.code\_execution -================================== - - -.. automodule:: autogen_ext.tools.code_execution - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.graphrag.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.graphrag.rst deleted file mode 100644 index d15d486b0c6b..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.graphrag.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.tools.graphrag -=========================== - - -.. automodule:: autogen_ext.tools.graphrag - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.http.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.http.rst deleted file mode 100644 index 7450fb9bfe08..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.http.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.tools.http -======================= - - -.. automodule:: autogen_ext.tools.http - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.langchain.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.langchain.rst deleted file mode 100644 index 416a3b550e6c..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.langchain.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.tools.langchain -============================ - - -.. automodule:: autogen_ext.tools.langchain - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.mcp.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.mcp.rst deleted file mode 100644 index 3c7ffd9085a5..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.mcp.rst +++ /dev/null @@ -1,8 +0,0 @@ -autogen\_ext.tools.mcp -====================== - - -.. automodule:: autogen_ext.tools.mcp - :members: - :undoc-members: - :show-inheritance: diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.semantic_kernel.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.semantic_kernel.rst deleted file mode 100644 index b3675fc1e1bb..000000000000 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_ext.tools.semantic_kernel.rst +++ /dev/null @@ -1,10 +0,0 @@ -autogen\_ext.tools.semantic_kernel -=================================== - - -.. automodule:: autogen_ext.tools.semantic_kernel - :members: - :undoc-members: - :show-inheritance: - - diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/custom-agents.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/custom-agents.ipynb deleted file mode 100644 index 1a526ee6106b..000000000000 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/custom-agents.ipynb +++ /dev/null @@ -1,739 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Custom Agents\n", - "\n", - "You may have agents with behaviors that do not fall into a preset. \n", - "In such cases, you can build custom agents.\n", - "\n", - "All agents in AgentChat inherit from {py:class}`~autogen_agentchat.agents.BaseChatAgent` \n", - "class and implement the following abstract methods and attributes:\n", - "\n", - "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages`: The abstract method that defines the behavior of the agent in response to messages. This method is called when the agent is asked to provide a response in {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run`. It returns a {py:class}`~autogen_agentchat.base.Response` object.\n", - "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_reset`: The abstract method that resets the agent to its initial state. This method is called when the agent is asked to reset itself.\n", - "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.produced_message_types`: The list of possible {py:class}`~autogen_agentchat.messages.ChatMessage` message types the agent can produce in its response.\n", - "\n", - "Optionally, you can implement the the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream` method to stream messages as they are generated by the agent. If this method is not implemented, the agent\n", - "uses the default implementation of {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream`\n", - "that calls the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages` method and\n", - "yields all messages in the response." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## CountDownAgent\n", - "\n", - "In this example, we create a simple agent that counts down from a given number to zero,\n", - "and produces a stream of messages with the current count." - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "3...\n", - "2...\n", - "1...\n", - "Done!\n" - ] - } - ], - "source": [ - "from typing import AsyncGenerator, List, Sequence\n", - "\n", - "from autogen_agentchat.agents import BaseChatAgent\n", - "from autogen_agentchat.base import Response\n", - "from autogen_agentchat.messages import AgentEvent, ChatMessage, TextMessage\n", - "from autogen_core import CancellationToken\n", - "\n", - "\n", - "class CountDownAgent(BaseChatAgent):\n", - " def __init__(self, name: str, count: int = 3):\n", - " super().__init__(name, \"A simple agent that counts down.\")\n", - " self._count = count\n", - "\n", - " @property\n", - " def produced_message_types(self) -> Sequence[type[ChatMessage]]:\n", - " return (TextMessage,)\n", - "\n", - " async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n", - " # Calls the on_messages_stream.\n", - " response: Response | None = None\n", - " async for message in self.on_messages_stream(messages, cancellation_token):\n", - " if isinstance(message, Response):\n", - " response = message\n", - " assert response is not None\n", - " return response\n", - "\n", - " async def on_messages_stream(\n", - " self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken\n", - " ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]:\n", - " inner_messages: List[AgentEvent | ChatMessage] = []\n", - " for i in range(self._count, 0, -1):\n", - " msg = TextMessage(content=f\"{i}...\", source=self.name)\n", - " inner_messages.append(msg)\n", - " yield msg\n", - " # The response is returned at the end of the stream.\n", - " # It contains the final message and all the inner messages.\n", - " yield Response(chat_message=TextMessage(content=\"Done!\", source=self.name), inner_messages=inner_messages)\n", - "\n", - " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", - " pass\n", - "\n", - "\n", - "async def run_countdown_agent() -> None:\n", - " # Create a countdown agent.\n", - " countdown_agent = CountDownAgent(\"countdown\")\n", - "\n", - " # Run the agent with a given task and stream the response.\n", - " async for message in countdown_agent.on_messages_stream([], CancellationToken()):\n", - " if isinstance(message, Response):\n", - " print(message.chat_message.content)\n", - " else:\n", - " print(message.content)\n", - "\n", - "\n", - "# Use asyncio.run(run_countdown_agent()) when running in a script.\n", - "await run_countdown_agent()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## ArithmeticAgent\n", - "\n", - "In this example, we create an agent class that can perform simple arithmetic operations\n", - "on a given integer. Then, we will use different instances of this agent class\n", - "in a {py:class}`~autogen_agentchat.teams.SelectorGroupChat`\n", - "to transform a given integer into another integer by applying a sequence of arithmetic operations.\n", - "\n", - "The `ArithmeticAgent` class takes an `operator_func` that takes an integer and returns an integer,\n", - "after applying an arithmetic operation to the integer.\n", - "In its `on_messages` method, it applies the `operator_func` to the integer in the input message,\n", - "and returns a response with the result." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from typing import Callable, Sequence\n", - "\n", - "from autogen_agentchat.agents import BaseChatAgent\n", - "from autogen_agentchat.base import Response\n", - "from autogen_agentchat.conditions import MaxMessageTermination\n", - "from autogen_agentchat.messages import ChatMessage\n", - "from autogen_agentchat.teams import SelectorGroupChat\n", - "from autogen_agentchat.ui import Console\n", - "from autogen_core import CancellationToken\n", - "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", - "\n", - "\n", - "class ArithmeticAgent(BaseChatAgent):\n", - " def __init__(self, name: str, description: str, operator_func: Callable[[int], int]) -> None:\n", - " super().__init__(name, description=description)\n", - " self._operator_func = operator_func\n", - " self._message_history: List[ChatMessage] = []\n", - "\n", - " @property\n", - " def produced_message_types(self) -> Sequence[type[ChatMessage]]:\n", - " return (TextMessage,)\n", - "\n", - " async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n", - " # Update the message history.\n", - " # NOTE: it is possible the messages is an empty list, which means the agent was selected previously.\n", - " self._message_history.extend(messages)\n", - " # Parse the number in the last message.\n", - " assert isinstance(self._message_history[-1], TextMessage)\n", - " number = int(self._message_history[-1].content)\n", - " # Apply the operator function to the number.\n", - " result = self._operator_func(number)\n", - " # Create a new message with the result.\n", - " response_message = TextMessage(content=str(result), source=self.name)\n", - " # Update the message history.\n", - " self._message_history.append(response_message)\n", - " # Return the response.\n", - " return Response(chat_message=response_message)\n", - "\n", - " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", - " pass" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{note}\n", - "The `on_messages` method may be called with an empty list of messages, in which\n", - "case it means the agent was called previously and is now being called again,\n", - "without any new messages from the caller. So it is important to keep a history\n", - "of the previous messages received by the agent, and use that history to generate\n", - "the response.\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now we can create a {py:class}`~autogen_agentchat.teams.SelectorGroupChat` with 5 instances of `ArithmeticAgent`:\n", - "\n", - "- one that adds 1 to the input integer,\n", - "- one that subtracts 1 from the input integer,\n", - "- one that multiplies the input integer by 2,\n", - "- one that divides the input integer by 2 and rounds down to the nearest integer, and\n", - "- one that returns the input integer unchanged.\n", - "\n", - "We then create a {py:class}`~autogen_agentchat.teams.SelectorGroupChat` with these agents,\n", - "and set the appropriate selector settings:\n", - "\n", - "- allow the same agent to be selected consecutively to allow for repeated operations, and\n", - "- customize the selector prompt to tailor the model's response to the specific task." - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "Apply the operations to turn the given number into 25.\n", - "---------- user ----------\n", - "10\n", - "---------- multiply_agent ----------\n", - "20\n", - "---------- add_agent ----------\n", - "21\n", - "---------- multiply_agent ----------\n", - "42\n", - "---------- divide_agent ----------\n", - "21\n", - "---------- add_agent ----------\n", - "22\n", - "---------- add_agent ----------\n", - "23\n", - "---------- add_agent ----------\n", - "24\n", - "---------- add_agent ----------\n", - "25\n", - "---------- Summary ----------\n", - "Number of messages: 10\n", - "Finish reason: Maximum number of messages 10 reached, current message count: 10\n", - "Total prompt tokens: 0\n", - "Total completion tokens: 0\n", - "Duration: 2.40 seconds\n" - ] - } - ], - "source": [ - "async def run_number_agents() -> None:\n", - " # Create agents for number operations.\n", - " add_agent = ArithmeticAgent(\"add_agent\", \"Adds 1 to the number.\", lambda x: x + 1)\n", - " multiply_agent = ArithmeticAgent(\"multiply_agent\", \"Multiplies the number by 2.\", lambda x: x * 2)\n", - " subtract_agent = ArithmeticAgent(\"subtract_agent\", \"Subtracts 1 from the number.\", lambda x: x - 1)\n", - " divide_agent = ArithmeticAgent(\"divide_agent\", \"Divides the number by 2 and rounds down.\", lambda x: x // 2)\n", - " identity_agent = ArithmeticAgent(\"identity_agent\", \"Returns the number as is.\", lambda x: x)\n", - "\n", - " # The termination condition is to stop after 10 messages.\n", - " termination_condition = MaxMessageTermination(10)\n", - "\n", - " # Create a selector group chat.\n", - " selector_group_chat = SelectorGroupChat(\n", - " [add_agent, multiply_agent, subtract_agent, divide_agent, identity_agent],\n", - " model_client=OpenAIChatCompletionClient(model=\"gpt-4o\"),\n", - " termination_condition=termination_condition,\n", - " allow_repeated_speaker=True, # Allow the same agent to speak multiple times, necessary for this task.\n", - " selector_prompt=(\n", - " \"Available roles:\\n{roles}\\nTheir job descriptions:\\n{participants}\\n\"\n", - " \"Current conversation history:\\n{history}\\n\"\n", - " \"Please select the most appropriate role for the next message, and only return the role name.\"\n", - " ),\n", - " )\n", - "\n", - " # Run the selector group chat with a given task and stream the response.\n", - " task: List[ChatMessage] = [\n", - " TextMessage(content=\"Apply the operations to turn the given number into 25.\", source=\"user\"),\n", - " TextMessage(content=\"10\", source=\"user\"),\n", - " ]\n", - " stream = selector_group_chat.run_stream(task=task)\n", - " await Console(stream)\n", - "\n", - "\n", - "# Use asyncio.run(run_number_agents()) when running in a script.\n", - "await run_number_agents()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From the output, we can see that the agents have successfully transformed the input integer\n", - "from 10 to 25 by choosing appropriate agents that apply the arithmetic operations in sequence." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Using Custom Model Clients in Custom Agents\n", - "\n", - "One of the key features of the {py:class}`~autogen_agentchat.agents.AssistantAgent` preset in AgentChat is that it takes a `model_client` argument and can use it in responding to messages. However, in some cases, you may want your agent to use a custom model client that is not currently supported (see [supported model clients](https://microsoft.github.io/autogen/dev/user-guide/core-user-guide/components/model-clients.html)) or custom model behaviours. \n", - "\n", - "You can accomplish this with a custom agent that implements *your custom model client*.\n", - "\n", - "In the example below, we will walk through an example of a custom agent that uses the [Google Gemini SDK](https://github.com/googleapis/python-genai) directly to respond to messages.\n", - "\n", - "> **Note:** You will need to install the [Google Gemini SDK](https://github.com/googleapis/python-genai) to run this example. You can install it using the following command: \n", - "\n", - "```bash\n", - "pip install google-genai\n", - "``` " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# !pip install google-genai\n", - "import os\n", - "from typing import AsyncGenerator, Sequence\n", - "\n", - "from autogen_agentchat.agents import BaseChatAgent\n", - "from autogen_agentchat.base import Response\n", - "from autogen_agentchat.messages import AgentEvent, ChatMessage\n", - "from autogen_core import CancellationToken\n", - "from autogen_core.model_context import UnboundedChatCompletionContext\n", - "from autogen_core.models import AssistantMessage, RequestUsage, UserMessage\n", - "from google import genai\n", - "from google.genai import types\n", - "\n", - "\n", - "class GeminiAssistantAgent(BaseChatAgent):\n", - " def __init__(\n", - " self,\n", - " name: str,\n", - " description: str = \"An agent that provides assistance with ability to use tools.\",\n", - " model: str = \"gemini-1.5-flash-002\",\n", - " api_key: str = os.environ[\"GEMINI_API_KEY\"],\n", - " system_message: str\n", - " | None = \"You are a helpful assistant that can respond to messages. Reply with TERMINATE when the task has been completed.\",\n", - " ):\n", - " super().__init__(name=name, description=description)\n", - " self._model_context = UnboundedChatCompletionContext()\n", - " self._model_client = genai.Client(api_key=api_key)\n", - " self._system_message = system_message\n", - " self._model = model\n", - "\n", - " @property\n", - " def produced_message_types(self) -> Sequence[type[ChatMessage]]:\n", - " return (TextMessage,)\n", - "\n", - " async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n", - " final_response = None\n", - " async for message in self.on_messages_stream(messages, cancellation_token):\n", - " if isinstance(message, Response):\n", - " final_response = message\n", - "\n", - " if final_response is None:\n", - " raise AssertionError(\"The stream should have returned the final result.\")\n", - "\n", - " return final_response\n", - "\n", - " async def on_messages_stream(\n", - " self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken\n", - " ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]:\n", - " # Add messages to the model context\n", - " for msg in messages:\n", - " await self._model_context.add_message(UserMessage(content=msg.content, source=msg.source))\n", - "\n", - " # Get conversation history\n", - " history = [\n", - " (msg.source if hasattr(msg, \"source\") else \"system\")\n", - " + \": \"\n", - " + (msg.content if isinstance(msg.content, str) else \"\")\n", - " + \"\\n\"\n", - " for msg in await self._model_context.get_messages()\n", - " ]\n", - " # Generate response using Gemini\n", - " response = self._model_client.models.generate_content(\n", - " model=self._model,\n", - " contents=f\"History: {history}\\nGiven the history, please provide a response\",\n", - " config=types.GenerateContentConfig(\n", - " system_instruction=self._system_message,\n", - " temperature=0.3,\n", - " ),\n", - " )\n", - "\n", - " # Create usage metadata\n", - " usage = RequestUsage(\n", - " prompt_tokens=response.usage_metadata.prompt_token_count,\n", - " completion_tokens=response.usage_metadata.candidates_token_count,\n", - " )\n", - "\n", - " # Add response to model context\n", - " await self._model_context.add_message(AssistantMessage(content=response.text, source=self.name))\n", - "\n", - " # Yield the final response\n", - " yield Response(\n", - " chat_message=TextMessage(content=response.text, source=self.name, models_usage=usage),\n", - " inner_messages=[],\n", - " )\n", - "\n", - " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", - " \"\"\"Reset the assistant by clearing the model context.\"\"\"\n", - " await self._model_context.clear()" - ] - }, - { - "cell_type": "code", - "execution_count": 38, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "What is the capital of New York?\n", - "---------- gemini_assistant ----------\n", - "Albany\n", - "TERMINATE\n", - "\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What is the capital of New York?', type='TextMessage'), TextMessage(source='gemini_assistant', models_usage=RequestUsage(prompt_tokens=46, completion_tokens=5), content='Albany\\nTERMINATE\\n', type='TextMessage')], stop_reason=None)" - ] - }, - "execution_count": 38, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "gemini_assistant = GeminiAssistantAgent(\"gemini_assistant\")\n", - "await Console(gemini_assistant.run_stream(task=\"What is the capital of New York?\"))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the example above, we have chosen to provide `model`, `api_key` and `system_message` as arguments - you can choose to provide any other arguments that are required by the model client you are using or fits with your application design. \n", - "\n", - "Now, let us explore how to use this custom agent as part of a team in AgentChat." - ] - }, - { - "cell_type": "code", - "execution_count": 39, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "Write a Haiku poem with 4 lines about the fall season.\n", - "---------- primary ----------\n", - "Crimson leaves cascade, \n", - "Whispering winds sing of change, \n", - "Chill wraps the fading, \n", - "Nature's quilt, rich and warm.\n", - "---------- gemini_critic ----------\n", - "The poem is good, but it has four lines instead of three. A haiku must have three lines with a 5-7-5 syllable structure. The content is evocative of autumn, but the form is incorrect. Please revise to adhere to the haiku's syllable structure.\n", - "\n", - "---------- primary ----------\n", - "Thank you for your feedback! Here’s a revised haiku that follows the 5-7-5 syllable structure:\n", - "\n", - "Crimson leaves drift down, \n", - "Chill winds whisper through the gold, \n", - "Autumn’s breath is near.\n", - "---------- gemini_critic ----------\n", - "The revised haiku is much improved. It correctly follows the 5-7-5 syllable structure and maintains the evocative imagery of autumn. APPROVE\n", - "\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Write a Haiku poem with 4 lines about the fall season.', type='TextMessage'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=33, completion_tokens=31), content=\"Crimson leaves cascade, \\nWhispering winds sing of change, \\nChill wraps the fading, \\nNature's quilt, rich and warm.\", type='TextMessage'), TextMessage(source='gemini_critic', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=60), content=\"The poem is good, but it has four lines instead of three. A haiku must have three lines with a 5-7-5 syllable structure. The content is evocative of autumn, but the form is incorrect. Please revise to adhere to the haiku's syllable structure.\\n\", type='TextMessage'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=141, completion_tokens=49), content='Thank you for your feedback! Here’s a revised haiku that follows the 5-7-5 syllable structure:\\n\\nCrimson leaves drift down, \\nChill winds whisper through the gold, \\nAutumn’s breath is near.', type='TextMessage'), TextMessage(source='gemini_critic', models_usage=RequestUsage(prompt_tokens=211, completion_tokens=32), content='The revised haiku is much improved. It correctly follows the 5-7-5 syllable structure and maintains the evocative imagery of autumn. APPROVE\\n', type='TextMessage')], stop_reason=\"Text 'APPROVE' mentioned\")" - ] - }, - "execution_count": 39, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "from autogen_agentchat.agents import AssistantAgent\n", - "from autogen_agentchat.conditions import TextMentionTermination\n", - "from autogen_agentchat.teams import RoundRobinGroupChat\n", - "from autogen_agentchat.ui import Console\n", - "\n", - "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\")\n", - "\n", - "# Create the primary agent.\n", - "primary_agent = AssistantAgent(\n", - " \"primary\",\n", - " model_client=model_client,\n", - " system_message=\"You are a helpful AI assistant.\",\n", - ")\n", - "\n", - "# Create a critic agent based on our new GeminiAssistantAgent.\n", - "gemini_critic_agent = GeminiAssistantAgent(\n", - " \"gemini_critic\",\n", - " system_message=\"Provide constructive feedback. Respond with 'APPROVE' to when your feedbacks are addressed.\",\n", - ")\n", - "\n", - "\n", - "# Define a termination condition that stops the task if the critic approves or after 10 messages.\n", - "termination = TextMentionTermination(\"APPROVE\") | MaxMessageTermination(10)\n", - "\n", - "# Create a team with the primary and critic agents.\n", - "team = RoundRobinGroupChat([primary_agent, gemini_critic_agent], termination_condition=termination)\n", - "\n", - "await Console(team.run_stream(task=\"Write a Haiku poem with 4 lines about the fall season.\"))\n", - "await model_client.close()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In section above, we show several very important concepts:\n", - "- We have developed a custom agent that uses the Google Gemini SDK to respond to messages. \n", - "- We show that this custom agent can be used as part of the broader AgentChat ecosystem - in this case as a participant in a {py:class}`~autogen_agentchat.teams.RoundRobinGroupChat` as long as it inherits from {py:class}`~autogen_agentchat.agents.BaseChatAgent`.\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Making the Custom Agent Declarative \n", - "\n", - "Autogen provides a [Component](https://microsoft.github.io/autogen/dev/user-guide/core-user-guide/framework/component-config.html) interface for making the configuration of components serializable to a declarative format. This is useful for saving and loading configurations, and for sharing configurations with others. \n", - "\n", - "We accomplish this by inheriting from the `Component` class and implementing the `_from_config` and `_to_config` methods.\n", - "The declarative class can be serialized to a JSON format using the `dump_component` method, and deserialized from a JSON format using the `load_component` method." - ] - }, - { - "cell_type": "code", - "execution_count": 40, - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "from typing import AsyncGenerator, Sequence\n", - "\n", - "from autogen_agentchat.agents import BaseChatAgent\n", - "from autogen_agentchat.base import Response\n", - "from autogen_agentchat.messages import AgentEvent, ChatMessage\n", - "from autogen_core import CancellationToken, Component\n", - "from pydantic import BaseModel\n", - "from typing_extensions import Self\n", - "\n", - "\n", - "class GeminiAssistantAgentConfig(BaseModel):\n", - " name: str\n", - " description: str = \"An agent that provides assistance with ability to use tools.\"\n", - " model: str = \"gemini-1.5-flash-002\"\n", - " system_message: str | None = None\n", - "\n", - "\n", - "class GeminiAssistantAgent(BaseChatAgent, Component[GeminiAssistantAgentConfig]): # type: ignore[no-redef]\n", - " component_config_schema = GeminiAssistantAgentConfig\n", - " # component_provider_override = \"mypackage.agents.GeminiAssistantAgent\"\n", - "\n", - " def __init__(\n", - " self,\n", - " name: str,\n", - " description: str = \"An agent that provides assistance with ability to use tools.\",\n", - " model: str = \"gemini-1.5-flash-002\",\n", - " api_key: str = os.environ[\"GEMINI_API_KEY\"],\n", - " system_message: str\n", - " | None = \"You are a helpful assistant that can respond to messages. Reply with TERMINATE when the task has been completed.\",\n", - " ):\n", - " super().__init__(name=name, description=description)\n", - " self._model_context = UnboundedChatCompletionContext()\n", - " self._model_client = genai.Client(api_key=api_key)\n", - " self._system_message = system_message\n", - " self._model = model\n", - "\n", - " @property\n", - " def produced_message_types(self) -> Sequence[type[ChatMessage]]:\n", - " return (TextMessage,)\n", - "\n", - " async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n", - " final_response = None\n", - " async for message in self.on_messages_stream(messages, cancellation_token):\n", - " if isinstance(message, Response):\n", - " final_response = message\n", - "\n", - " if final_response is None:\n", - " raise AssertionError(\"The stream should have returned the final result.\")\n", - "\n", - " return final_response\n", - "\n", - " async def on_messages_stream(\n", - " self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken\n", - " ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]:\n", - " # Add messages to the model context\n", - " for msg in messages:\n", - " await self._model_context.add_message(UserMessage(content=msg.content, source=msg.source))\n", - "\n", - " # Get conversation history\n", - " history = [\n", - " (msg.source if hasattr(msg, \"source\") else \"system\")\n", - " + \": \"\n", - " + (msg.content if isinstance(msg.content, str) else \"\")\n", - " + \"\\n\"\n", - " for msg in await self._model_context.get_messages()\n", - " ]\n", - "\n", - " # Generate response using Gemini\n", - " response = self._model_client.models.generate_content(\n", - " model=self._model,\n", - " contents=f\"History: {history}\\nGiven the history, please provide a response\",\n", - " config=types.GenerateContentConfig(\n", - " system_instruction=self._system_message,\n", - " temperature=0.3,\n", - " ),\n", - " )\n", - "\n", - " # Create usage metadata\n", - " usage = RequestUsage(\n", - " prompt_tokens=response.usage_metadata.prompt_token_count,\n", - " completion_tokens=response.usage_metadata.candidates_token_count,\n", - " )\n", - "\n", - " # Add response to model context\n", - " await self._model_context.add_message(AssistantMessage(content=response.text, source=self.name))\n", - "\n", - " # Yield the final response\n", - " yield Response(\n", - " chat_message=TextMessage(content=response.text, source=self.name, models_usage=usage),\n", - " inner_messages=[],\n", - " )\n", - "\n", - " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", - " \"\"\"Reset the assistant by clearing the model context.\"\"\"\n", - " await self._model_context.clear()\n", - "\n", - " @classmethod\n", - " def _from_config(cls, config: GeminiAssistantAgentConfig) -> Self:\n", - " return cls(\n", - " name=config.name, description=config.description, model=config.model, system_message=config.system_message\n", - " )\n", - "\n", - " def _to_config(self) -> GeminiAssistantAgentConfig:\n", - " return GeminiAssistantAgentConfig(\n", - " name=self.name,\n", - " description=self.description,\n", - " model=self._model,\n", - " system_message=self._system_message,\n", - " )" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now that we have the required methods implemented, we can now load and dump the custom agent to and from a JSON format, and then load the agent from the JSON format.\n", - " \n", - " > Note: You should set the `component_provider_override` class variable to the full path of the module containing the custom agent class e.g., (`mypackage.agents.GeminiAssistantAgent`). This is used by `load_component` method to determine how to instantiate the class. \n", - " " - ] - }, - { - "cell_type": "code", - "execution_count": 41, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "{\n", - " \"provider\": \"__main__.GeminiAssistantAgent\",\n", - " \"component_type\": \"agent\",\n", - " \"version\": 1,\n", - " \"component_version\": 1,\n", - " \"description\": null,\n", - " \"label\": \"GeminiAssistantAgent\",\n", - " \"config\": {\n", - " \"name\": \"gemini_assistant\",\n", - " \"description\": \"An agent that provides assistance with ability to use tools.\",\n", - " \"model\": \"gemini-1.5-flash-002\",\n", - " \"system_message\": \"You are a helpful assistant that can respond to messages. Reply with TERMINATE when the task has been completed.\"\n", - " }\n", - "}\n", - "<__main__.GeminiAssistantAgent object at 0x11a5c5a90>\n" - ] - } - ], - "source": [ - "gemini_assistant = GeminiAssistantAgent(\"gemini_assistant\")\n", - "config = gemini_assistant.dump_component()\n", - "print(config.model_dump_json(indent=2))\n", - "loaded_agent = GeminiAssistantAgent.load_component(config)\n", - "print(loaded_agent)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Next Steps \n", - "\n", - "So far, we have seen how to create custom agents, add custom model clients to agents, and make custom agents declarative. There are a few ways in which this basic sample can be extended:\n", - "\n", - "- Extend the Gemini model client to handle function calling similar to the {py:class}`~autogen_agentchat.agents.AssistantAgent` class. https://ai.google.dev/gemini-api/docs/function-calling \n", - "- Implement a package with a custom agent and experiment with using its declarative format in a tool like [AutoGen Studio](https://microsoft.github.io/autogen/stable/user-guide/autogenstudio-user-guide/index.html)." - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.9" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/jaeger.png b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/jaeger.png deleted file mode 100644 index ccfbe9301f33..000000000000 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/jaeger.png +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2375682ac24d72954a16cecee759d024e969f37fa6f9514bfb736873a540ea4c -size 283931 diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/memory.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/memory.ipynb deleted file mode 100644 index ca6f3e7c6119..000000000000 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/memory.ipynb +++ /dev/null @@ -1,442 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Memory and RAG\n", - "\n", - "There are several use cases where it is valuable to maintain a _store_ of useful facts that can be intelligently added to the context of the agent just before a specific step. The typically use case here is a RAG pattern where a query is used to retrieve relevant information from a database that is then added to the agent's context.\n", - "\n", - "\n", - "AgentChat provides a {py:class}`~autogen_core.memory.Memory` protocol that can be extended to provide this functionality. The key methods are `query`, `update_context`, `add`, `clear`, and `close`. \n", - "\n", - "- `add`: add new entries to the memory store\n", - "- `query`: retrieve relevant information from the memory store \n", - "- `update_context`: mutate an agent's internal `model_context` by adding the retrieved information (used in the {py:class}`~autogen_agentchat.agents.AssistantAgent` class) \n", - "- `clear`: clear all entries from the memory store\n", - "- `close`: clean up any resources used by the memory store \n", - "\n", - "\n", - "## ListMemory Example\n", - "\n", - "{py:class}~autogen_core.memory.ListMemory is provided as an example implementation of the {py:class}~autogen_core.memory.Memory protocol. It is a simple list-based memory implementation that maintains memories in chronological order, appending the most recent memories to the model's context. The implementation is designed to be straightforward and predictable, making it easy to understand and debug.\n", - "In the following example, we will use ListMemory to maintain a memory bank of user preferences and demonstrate how it can be used to provide consistent context for agent responses over time." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from autogen_agentchat.agents import AssistantAgent\n", - "from autogen_agentchat.ui import Console\n", - "from autogen_core.memory import ListMemory, MemoryContent, MemoryMimeType\n", - "from autogen_ext.models.openai import OpenAIChatCompletionClient" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Initialize user memory\n", - "user_memory = ListMemory()\n", - "\n", - "# Add user preferences to memory\n", - "await user_memory.add(MemoryContent(content=\"The weather should be in metric units\", mime_type=MemoryMimeType.TEXT))\n", - "\n", - "await user_memory.add(MemoryContent(content=\"Meal recipe must be vegan\", mime_type=MemoryMimeType.TEXT))\n", - "\n", - "\n", - "async def get_weather(city: str, units: str = \"imperial\") -> str:\n", - " if units == \"imperial\":\n", - " return f\"The weather in {city} is 73 °F and Sunny.\"\n", - " elif units == \"metric\":\n", - " return f\"The weather in {city} is 23 °C and Sunny.\"\n", - " else:\n", - " return f\"Sorry, I don't know the weather in {city}.\"\n", - "\n", - "\n", - "assistant_agent = AssistantAgent(\n", - " name=\"assistant_agent\",\n", - " model_client=OpenAIChatCompletionClient(\n", - " model=\"gpt-4o-2024-08-06\",\n", - " ),\n", - " tools=[get_weather],\n", - " memory=[user_memory],\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Run the agent with a task.\n", - "stream = assistant_agent.run_stream(task=\"What is the weather in New York?\")\n", - "await Console(stream)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can inspect that the `assistant_agent` model_context is actually updated with the retrieved memory entries. The `transform` method is used to format the retrieved memory entries into a string that can be used by the agent. In this case, we simply concatenate the content of each memory entry into a single string." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "await assistant_agent._model_context.get_messages()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We see above that the weather is returned in Centigrade as stated in the user preferences. \n", - "\n", - "Similarly, assuming we ask a separate question about generating a meal plan, the agent is able to retrieve relevant information from the memory store and provide a personalized (vegan) response." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "stream = assistant_agent.run_stream(task=\"Write brief meal recipe with broth\")\n", - "await Console(stream)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Custom Memory Stores (Vector DBs, etc.)\n", - "\n", - "You can build on the `Memory` protocol to implement more complex memory stores. For example, you could implement a custom memory store that uses a vector database to store and retrieve information, or a memory store that uses a machine learning model to generate personalized responses based on the user's preferences etc.\n", - "\n", - "Specifically, you will need to overload the `add`, `query` and `update_context` methods to implement the desired functionality and pass the memory store to your agent.\n", - "\n", - "\n", - "Currently the following example memory stores are available as part of the {py:class}`~autogen_ext` extensions package. \n", - "\n", - "- `autogen_ext.memory.chromadb.ChromaDBVectorMemory`: A memory store that uses a vector database to store and retrieve information. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "from pathlib import Path\n", - "\n", - "from autogen_agentchat.agents import AssistantAgent\n", - "from autogen_agentchat.ui import Console\n", - "from autogen_core.memory import MemoryContent, MemoryMimeType\n", - "from autogen_ext.memory.chromadb import ChromaDBVectorMemory, PersistentChromaDBVectorMemoryConfig\n", - "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", - "\n", - "# Initialize ChromaDB memory with custom config\n", - "chroma_user_memory = ChromaDBVectorMemory(\n", - " config=PersistentChromaDBVectorMemoryConfig(\n", - " collection_name=\"preferences\",\n", - " persistence_path=os.path.join(str(Path.home()), \".chromadb_autogen\"),\n", - " k=2, # Return top k results\n", - " score_threshold=0.4, # Minimum similarity score\n", - " )\n", - ")\n", - "# a HttpChromaDBVectorMemoryConfig is also supported for connecting to a remote ChromaDB server\n", - "\n", - "# Add user preferences to memory\n", - "await chroma_user_memory.add(\n", - " MemoryContent(\n", - " content=\"The weather should be in metric units\",\n", - " mime_type=MemoryMimeType.TEXT,\n", - " metadata={\"category\": \"preferences\", \"type\": \"units\"},\n", - " )\n", - ")\n", - "\n", - "await chroma_user_memory.add(\n", - " MemoryContent(\n", - " content=\"Meal recipe must be vegan\",\n", - " mime_type=MemoryMimeType.TEXT,\n", - " metadata={\"category\": \"preferences\", \"type\": \"dietary\"},\n", - " )\n", - ")\n", - "\n", - "model_client = OpenAIChatCompletionClient(\n", - " model=\"gpt-4o\",\n", - ")\n", - "\n", - "# Create assistant agent with ChromaDB memory\n", - "assistant_agent = AssistantAgent(\n", - " name=\"assistant_agent\",\n", - " model_client=model_client,\n", - " tools=[get_weather],\n", - " memory=[chroma_user_memory],\n", - ")\n", - "\n", - "stream = assistant_agent.run_stream(task=\"What is the weather in New York?\")\n", - "await Console(stream)\n", - "\n", - "await model_client.close()\n", - "await chroma_user_memory.close()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note that you can also serialize the ChromaDBVectorMemory and save it to disk." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "chroma_user_memory.dump_component().model_dump_json()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## RAG Agent: Putting It All Together\n", - "\n", - "The RAG (Retrieval Augmented Generation) pattern which is common in building AI systems encompasses two distinct phases:\n", - "\n", - "1. **Indexing**: Loading documents, chunking them, and storing them in a vector database\n", - "2. **Retrieval**: Finding and using relevant chunks during conversation runtime\n", - "\n", - "In our previous examples, we manually added items to memory and passed them to our agents. In practice, the indexing process is usually automated and based on much larger document sources like product documentation, internal files, or knowledge bases.\n", - "\n", - "> Note: The quality of a RAG system is dependent on the quality of the chunking and retrieval process (models, embeddings, etc.). You may need to experiement with more advanced chunking and retrieval models to get the best results.\n", - "\n", - "### Building a Simple RAG Agent\n", - "\n", - "To begin, let's create a simple document indexer that we will used to load documents, chunk them, and store them in a `ChromaDBVectorMemory` memory store. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import re\n", - "from typing import List\n", - "\n", - "import aiofiles\n", - "import aiohttp\n", - "from autogen_core.memory import Memory, MemoryContent, MemoryMimeType\n", - "\n", - "\n", - "class SimpleDocumentIndexer:\n", - " \"\"\"Basic document indexer for AutoGen Memory.\"\"\"\n", - "\n", - " def __init__(self, memory: Memory, chunk_size: int = 1500) -> None:\n", - " self.memory = memory\n", - " self.chunk_size = chunk_size\n", - "\n", - " async def _fetch_content(self, source: str) -> str:\n", - " \"\"\"Fetch content from URL or file.\"\"\"\n", - " if source.startswith((\"http://\", \"https://\")):\n", - " async with aiohttp.ClientSession() as session:\n", - " async with session.get(source) as response:\n", - " return await response.text()\n", - " else:\n", - " async with aiofiles.open(source, \"r\", encoding=\"utf-8\") as f:\n", - " return await f.read()\n", - "\n", - " def _strip_html(self, text: str) -> str:\n", - " \"\"\"Remove HTML tags and normalize whitespace.\"\"\"\n", - " text = re.sub(r\"<[^>]*>\", \" \", text)\n", - " text = re.sub(r\"\\s+\", \" \", text)\n", - " return text.strip()\n", - "\n", - " def _split_text(self, text: str) -> List[str]:\n", - " \"\"\"Split text into fixed-size chunks.\"\"\"\n", - " chunks: list[str] = []\n", - " # Just split text into fixed-size chunks\n", - " for i in range(0, len(text), self.chunk_size):\n", - " chunk = text[i : i + self.chunk_size]\n", - " chunks.append(chunk.strip())\n", - " return chunks\n", - "\n", - " async def index_documents(self, sources: List[str]) -> int:\n", - " \"\"\"Index documents into memory.\"\"\"\n", - " total_chunks = 0\n", - "\n", - " for source in sources:\n", - " try:\n", - " content = await self._fetch_content(source)\n", - "\n", - " # Strip HTML if content appears to be HTML\n", - " if \"<\" in content and \">\" in content:\n", - " content = self._strip_html(content)\n", - "\n", - " chunks = self._split_text(content)\n", - "\n", - " for i, chunk in enumerate(chunks):\n", - " await self.memory.add(\n", - " MemoryContent(\n", - " content=chunk, mime_type=MemoryMimeType.TEXT, metadata={\"source\": source, \"chunk_index\": i}\n", - " )\n", - " )\n", - "\n", - " total_chunks += len(chunks)\n", - "\n", - " except Exception as e:\n", - " print(f\"Error indexing {source}: {str(e)}\")\n", - "\n", - " return total_chunks" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "\n", - " \n", - "Now let's use our indexer with ChromaDBVectorMemory to build a complete RAG agent:\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Indexed 72 chunks from 4 AutoGen documents\n" - ] - } - ], - "source": [ - "import os\n", - "from pathlib import Path\n", - "\n", - "from autogen_agentchat.agents import AssistantAgent\n", - "from autogen_agentchat.ui import Console\n", - "from autogen_ext.memory.chromadb import ChromaDBVectorMemory, PersistentChromaDBVectorMemoryConfig\n", - "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", - "\n", - "# Initialize vector memory\n", - "\n", - "rag_memory = ChromaDBVectorMemory(\n", - " config=PersistentChromaDBVectorMemoryConfig(\n", - " collection_name=\"autogen_docs\",\n", - " persistence_path=os.path.join(str(Path.home()), \".chromadb_autogen\"),\n", - " k=3, # Return top 3 results\n", - " score_threshold=0.4, # Minimum similarity score\n", - " )\n", - ")\n", - "\n", - "await rag_memory.clear() # Clear existing memory\n", - "\n", - "\n", - "# Index AutoGen documentation\n", - "async def index_autogen_docs() -> None:\n", - " indexer = SimpleDocumentIndexer(memory=rag_memory)\n", - " sources = [\n", - " \"https://raw.githubusercontent.com/microsoft/autogen/main/README.md\",\n", - " \"https://microsoft.github.io/autogen/dev/user-guide/agentchat-user-guide/tutorial/agents.html\",\n", - " \"https://microsoft.github.io/autogen/dev/user-guide/agentchat-user-guide/tutorial/teams.html\",\n", - " \"https://microsoft.github.io/autogen/dev/user-guide/agentchat-user-guide/tutorial/termination.html\",\n", - " ]\n", - " chunks: int = await indexer.index_documents(sources)\n", - " print(f\"Indexed {chunks} chunks from {len(sources)} AutoGen documents\")\n", - "\n", - "\n", - "await index_autogen_docs()" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "What is AgentChat?\n", - "Query results: results=[MemoryContent(content='ng OpenAI\\'s GPT-4o model. See [other supported models](https://microsoft.github.io/autogen/stable/user-guide/agentchat-user-guide/tutorial/models.html). ```python import asyncio from autogen_agentchat.agents import AssistantAgent from autogen_ext.models.openai import OpenAIChatCompletionClient async def main() -> None: model_client = OpenAIChatCompletionClient(model=\"gpt-4o\") agent = AssistantAgent(\"assistant\", model_client=model_client) print(await agent.run(task=\"Say \\'Hello World!\\'\")) await model_client.close() asyncio.run(main()) ``` ### Web Browsing Agent Team Create a group chat team with a web surfer agent and a user proxy agent for web browsing tasks. You need to install [playwright](https://playwright.dev/python/docs/library). ```python # pip install -U autogen-agentchat autogen-ext[openai,web-surfer] # playwright install import asyncio from autogen_agentchat.agents import UserProxyAgent from autogen_agentchat.conditions import TextMentionTermination from autogen_agentchat.teams import RoundRobinGroupChat from autogen_agentchat.ui import Console from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.agents.web_surfer import MultimodalWebSurfer async def main() -> None: model_client = OpenAIChatCompletionClient(model=\"gpt-4o\") # The web surfer will open a Chromium browser window to perform web browsing tasks. web_surfer = MultimodalWebSurfer(\"web_surfer\", model_client, headless=False, animate_actions=True) # The user proxy agent is used to ge', mime_type='MemoryMimeType.TEXT', metadata={'chunk_index': 1, 'mime_type': 'MemoryMimeType.TEXT', 'source': 'https://raw.githubusercontent.com/microsoft/autogen/main/README.md', 'score': 0.48810458183288574, 'id': '16088e03-0153-4da3-9dec-643b39c549f5'}), MemoryContent(content='els_usage=None content='AutoGen is a programming framework for building multi-agent applications.' type='ToolCallSummaryMessage' The call to the on_messages() method returns a Response that contains the agent’s final response in the chat_message attribute, as well as a list of inner messages in the inner_messages attribute, which stores the agent’s “thought process” that led to the final response. Note It is important to note that on_messages() will update the internal state of the agent – it will add the messages to the agent’s history. So you should call this method with new messages. You should not repeatedly call this method with the same messages or the complete history. Note Unlike in v0.2 AgentChat, the tools are executed by the same agent directly within the same call to on_messages() . By default, the agent will return the result of the tool call as the final response. You can also call the run() method, which is a convenience method that calls on_messages() . It follows the same interface as Teams and returns a TaskResult object. Multi-Modal Input # The AssistantAgent can handle multi-modal input by providing the input as a MultiModalMessage . from io import BytesIO import PIL import requests from autogen_agentchat.messages import MultiModalMessage from autogen_core import Image # Create a multi-modal message with random image and text. pil_image = PIL . Image . open ( BytesIO ( requests . get ( "https://picsum.photos/300/200" ) . content )', mime_type='MemoryMimeType.TEXT', metadata={'chunk_index': 3, 'mime_type': 'MemoryMimeType.TEXT', 'source': 'https://microsoft.github.io/autogen/dev/user-guide/agentchat-user-guide/tutorial/agents.html', 'score': 0.4665141701698303, 'id': '3d603b62-7cab-4f74-b671-586fe36306f2'}), MemoryContent(content='AgentChat Termination Termination # In the previous section, we explored how to define agents, and organize them into teams that can solve tasks. However, a run can go on forever, and in many cases, we need to know when to stop them. This is the role of the termination condition. AgentChat supports several termination condition by providing a base TerminationCondition class and several implementations that inherit from it. A termination condition is a callable that takes a sequence of AgentEvent or ChatMessage objects since the last time the condition was called , and returns a StopMessage if the conversation should be terminated, or None otherwise. Once a termination condition has been reached, it must be reset by calling reset() before it can be used again. Some important things to note about termination conditions: They are stateful but reset automatically after each run ( run() or run_stream() ) is finished. They can be combined using the AND and OR operators. Note For group chat teams (i.e., RoundRobinGroupChat , SelectorGroupChat , and Swarm ), the termination condition is called after each agent responds. While a response may contain multiple inner messages, the team calls its termination condition just once for all the messages from a single response. So the condition is called with the “delta sequence” of messages since the last time it was called. Built-In Termination Conditions: MaxMessageTermination : Stops after a specified number of messages have been produced,', mime_type='MemoryMimeType.TEXT', metadata={'chunk_index': 1, 'mime_type': 'MemoryMimeType.TEXT', 'source': 'https://microsoft.github.io/autogen/dev/user-guide/agentchat-user-guide/tutorial/termination.html', 'score': 0.461774212772051, 'id': '699ef490-d108-4cd3-b629-c1198d6b78ba'})]\n", - "---------- rag_assistant ----------\n", - "[MemoryContent(content='ng OpenAI\\'s GPT-4o model. See [other supported models](https://microsoft.github.io/autogen/stable/user-guide/agentchat-user-guide/tutorial/models.html). ```python import asyncio from autogen_agentchat.agents import AssistantAgent from autogen_ext.models.openai import OpenAIChatCompletionClient async def main() -> None: model_client = OpenAIChatCompletionClient(model=\"gpt-4o\") agent = AssistantAgent(\"assistant\", model_client=model_client) print(await agent.run(task=\"Say \\'Hello World!\\'\")) await model_client.close() asyncio.run(main()) ``` ### Web Browsing Agent Team Create a group chat team with a web surfer agent and a user proxy agent for web browsing tasks. You need to install [playwright](https://playwright.dev/python/docs/library). ```python # pip install -U autogen-agentchat autogen-ext[openai,web-surfer] # playwright install import asyncio from autogen_agentchat.agents import UserProxyAgent from autogen_agentchat.conditions import TextMentionTermination from autogen_agentchat.teams import RoundRobinGroupChat from autogen_agentchat.ui import Console from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.agents.web_surfer import MultimodalWebSurfer async def main() -> None: model_client = OpenAIChatCompletionClient(model=\"gpt-4o\") # The web surfer will open a Chromium browser window to perform web browsing tasks. web_surfer = MultimodalWebSurfer(\"web_surfer\", model_client, headless=False, animate_actions=True) # The user proxy agent is used to ge', mime_type='MemoryMimeType.TEXT', metadata={'chunk_index': 1, 'mime_type': 'MemoryMimeType.TEXT', 'source': 'https://raw.githubusercontent.com/microsoft/autogen/main/README.md', 'score': 0.48810458183288574, 'id': '16088e03-0153-4da3-9dec-643b39c549f5'}), MemoryContent(content='els_usage=None content='AutoGen is a programming framework for building multi-agent applications.' type='ToolCallSummaryMessage' The call to the on_messages() method returns a Response that contains the agent’s final response in the chat_message attribute, as well as a list of inner messages in the inner_messages attribute, which stores the agent’s “thought process” that led to the final response. Note It is important to note that on_messages() will update the internal state of the agent – it will add the messages to the agent’s history. So you should call this method with new messages. You should not repeatedly call this method with the same messages or the complete history. Note Unlike in v0.2 AgentChat, the tools are executed by the same agent directly within the same call to on_messages() . By default, the agent will return the result of the tool call as the final response. You can also call the run() method, which is a convenience method that calls on_messages() . It follows the same interface as Teams and returns a TaskResult object. Multi-Modal Input # The AssistantAgent can handle multi-modal input by providing the input as a MultiModalMessage . from io import BytesIO import PIL import requests from autogen_agentchat.messages import MultiModalMessage from autogen_core import Image # Create a multi-modal message with random image and text. pil_image = PIL . Image . open ( BytesIO ( requests . get ( "https://picsum.photos/300/200" ) . content )', mime_type='MemoryMimeType.TEXT', metadata={'chunk_index': 3, 'mime_type': 'MemoryMimeType.TEXT', 'source': 'https://microsoft.github.io/autogen/dev/user-guide/agentchat-user-guide/tutorial/agents.html', 'score': 0.4665141701698303, 'id': '3d603b62-7cab-4f74-b671-586fe36306f2'}), MemoryContent(content='AgentChat Termination Termination # In the previous section, we explored how to define agents, and organize them into teams that can solve tasks. However, a run can go on forever, and in many cases, we need to know when to stop them. This is the role of the termination condition. AgentChat supports several termination condition by providing a base TerminationCondition class and several implementations that inherit from it. A termination condition is a callable that takes a sequence of AgentEvent or ChatMessage objects since the last time the condition was called , and returns a StopMessage if the conversation should be terminated, or None otherwise. Once a termination condition has been reached, it must be reset by calling reset() before it can be used again. Some important things to note about termination conditions: They are stateful but reset automatically after each run ( run() or run_stream() ) is finished. They can be combined using the AND and OR operators. Note For group chat teams (i.e., RoundRobinGroupChat , SelectorGroupChat , and Swarm ), the termination condition is called after each agent responds. While a response may contain multiple inner messages, the team calls its termination condition just once for all the messages from a single response. So the condition is called with the “delta sequence” of messages since the last time it was called. Built-In Termination Conditions: MaxMessageTermination : Stops after a specified number of messages have been produced,', mime_type='MemoryMimeType.TEXT', metadata={'chunk_index': 1, 'mime_type': 'MemoryMimeType.TEXT', 'source': 'https://microsoft.github.io/autogen/dev/user-guide/agentchat-user-guide/tutorial/termination.html', 'score': 0.461774212772051, 'id': '699ef490-d108-4cd3-b629-c1198d6b78ba'})]\n", - "---------- rag_assistant ----------\n", - "AgentChat is part of the AutoGen framework, a programming environment for building multi-agent applications. In AgentChat, agents can interact with each other and with users to perform various tasks, including web browsing and engaging in dialogue. It utilizes models from OpenAI for chat completions and supports multi-modal input, which means agents can handle inputs that include both text and images. Additionally, AgentChat provides mechanisms to define termination conditions to control when a conversation or task should be concluded, ensuring that the agent interactions are efficient and goal-oriented. TERMINATE\n" - ] - } - ], - "source": [ - "# Create our RAG assistant agent\n", - "rag_assistant = AssistantAgent(\n", - " name=\"rag_assistant\", model_client=OpenAIChatCompletionClient(model=\"gpt-4o\"), memory=[rag_memory]\n", - ")\n", - "\n", - "# Ask questions about AutoGen\n", - "stream = rag_assistant.run_stream(task=\"What is AgentChat?\")\n", - "await Console(stream)\n", - "\n", - "# Remember to close the memory when done\n", - "await rag_memory.close()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This implementation provides a RAG agent that can answer questions based on AutoGen documentation. When a question is asked, the Memory system retrieves relevant chunks and adds them to the context, enabling the assistant to generate informed responses.\n", - "\n", - "For production systems, you might want to:\n", - "1. Implement more sophisticated chunking strategies\n", - "2. Add metadata filtering capabilities\n", - "3. Customize the retrieval scoring\n", - "4. Optimize embedding models for your specific domain\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.9" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/selector-group-chat.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/selector-group-chat.ipynb deleted file mode 100644 index 522c26b2098c..000000000000 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/selector-group-chat.ipynb +++ /dev/null @@ -1,1026 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Selector Group Chat" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "{py:class}`~autogen_agentchat.teams.SelectorGroupChat` implements a team where participants take turns broadcasting messages to all other members. A generative model (e.g., an LLM) selects the next speaker based on the shared context, enabling dynamic, context-aware collaboration.\n", - "\n", - "Key features include:\n", - "\n", - "- Model-based speaker selection\n", - "- Configurable participant roles and descriptions\n", - "- Prevention of consecutive turns by the same speaker (optional)\n", - "- Customizable selection prompting\n", - "- Customizable selection function to override the default model-based selection\n", - "- Customizable candidate function to narrow-down the set of agents for selection using model\n", - "\n", - "```{note}\n", - "{py:class}`~autogen_agentchat.teams.SelectorGroupChat` is a high-level API. For more control and customization, refer to the [Group Chat Pattern](../core-user-guide/design-patterns/group-chat.ipynb) in the Core API documentation to implement your own group chat logic.\n", - "```\n", - "\n", - "## How Does it Work?\n", - "\n", - "{py:class}`~autogen_agentchat.teams.SelectorGroupChat` is a group chat similar to {py:class}`~autogen_agentchat.teams.RoundRobinGroupChat`,\n", - "but with a model-based next speaker selection mechanism.\n", - "When the team receives a task through {py:meth}`~autogen_agentchat.teams.BaseGroupChat.run` or {py:meth}`~autogen_agentchat.teams.BaseGroupChat.run_stream`,\n", - "the following steps are executed:\n", - "\n", - "1. The team analyzes the current conversation context, including the conversation history and participants' {py:attr}`~autogen_agentchat.base.ChatAgent.name` and {py:attr}`~autogen_agentchat.base.ChatAgent.description` attributes, to determine the next speaker using a model. By default, the team will not select the same speak consecutively unless it is the only agent available. This can be changed by setting `allow_repeated_speaker=True`. You can also override the model by providing a custom selection function.\n", - "2. The team prompts the selected speaker agent to provide a response, which is then **broadcasted** to all other participants.\n", - "3. The termination condition is checked to determine if the conversation should end, if not, the process repeats from step 1.\n", - "4. When the conversation ends, the team returns the {py:class}`~autogen_agentchat.base.TaskResult` containing the conversation history from this task.\n", - "\n", - "Once the team finishes the task, the conversation context is kept within the team and all participants, so the next task can continue from the previous conversation context.\n", - "You can reset the conversation context by calling {py:meth}`~autogen_agentchat.teams.BaseGroupChat.reset`.\n", - "\n", - "In this section, we will demonstrate how to use {py:class}`~autogen_agentchat.teams.SelectorGroupChat` with a simple example for a web search and data analysis task." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Example: Web Search/Analysis" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [], - "source": [ - "from typing import List, Sequence\n", - "\n", - "from autogen_agentchat.agents import AssistantAgent, UserProxyAgent\n", - "from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination\n", - "from autogen_agentchat.messages import AgentEvent, ChatMessage\n", - "from autogen_agentchat.teams import SelectorGroupChat\n", - "from autogen_agentchat.ui import Console\n", - "from autogen_ext.models.openai import OpenAIChatCompletionClient" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Agents\n", - "\n", - "![Selector Group Chat](selector-group-chat.svg)\n", - "\n", - "This system uses three specialized agents:\n", - "\n", - "- **Planning Agent**: The strategic coordinator that breaks down complex tasks into manageable subtasks. \n", - "- **Web Search Agent**: An information retrieval specialist that interfaces with the `search_web_tool`.\n", - "- **Data Analyst Agent**: An agent specialist in performing calculations equipped with `percentage_change_tool`. " - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The tools `search_web_tool` and `percentage_change_tool` are external tools that the agents can use to perform their tasks." - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": {}, - "outputs": [], - "source": [ - "# Note: This example uses mock tools instead of real APIs for demonstration purposes\n", - "def search_web_tool(query: str) -> str:\n", - " if \"2006-2007\" in query:\n", - " return \"\"\"Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", - " Udonis Haslem: 844 points\n", - " Dwayne Wade: 1397 points\n", - " James Posey: 550 points\n", - " ...\n", - " \"\"\"\n", - " elif \"2007-2008\" in query:\n", - " return \"The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\"\n", - " elif \"2008-2009\" in query:\n", - " return \"The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\"\n", - " return \"No data found.\"\n", - "\n", - "\n", - "def percentage_change_tool(start: float, end: float) -> float:\n", - " return ((end - start) / start) * 100" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Let's create the specialized agents using the {py:class}`~autogen_agentchat.agents.AssistantAgent` class.\n", - "It is important to note that the agents' {py:attr}`~autogen_agentchat.base.ChatAgent.name` and {py:attr}`~autogen_agentchat.base.ChatAgent.description` attributes are used by the model to determine the next speaker,\n", - "so it is recommended to provide meaningful names and descriptions." - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": {}, - "outputs": [], - "source": [ - "model_client = OpenAIChatCompletionClient(model=\"gpt-4o\")\n", - "\n", - "planning_agent = AssistantAgent(\n", - " \"PlanningAgent\",\n", - " description=\"An agent for planning tasks, this agent should be the first to engage when given a new task.\",\n", - " model_client=model_client,\n", - " system_message=\"\"\"\n", - " You are a planning agent.\n", - " Your job is to break down complex tasks into smaller, manageable subtasks.\n", - " Your team members are:\n", - " WebSearchAgent: Searches for information\n", - " DataAnalystAgent: Performs calculations\n", - "\n", - " You only plan and delegate tasks - you do not execute them yourself.\n", - "\n", - " When assigning tasks, use this format:\n", - " 1. : \n", - "\n", - " After all tasks are complete, summarize the findings and end with \"TERMINATE\".\n", - " \"\"\",\n", - ")\n", - "\n", - "web_search_agent = AssistantAgent(\n", - " \"WebSearchAgent\",\n", - " description=\"An agent for searching information on the web.\",\n", - " tools=[search_web_tool],\n", - " model_client=model_client,\n", - " system_message=\"\"\"\n", - " You are a web search agent.\n", - " Your only tool is search_tool - use it to find information.\n", - " You make only one search call at a time.\n", - " Once you have the results, you never do calculations based on them.\n", - " \"\"\",\n", - ")\n", - "\n", - "data_analyst_agent = AssistantAgent(\n", - " \"DataAnalystAgent\",\n", - " description=\"An agent for performing calculations.\",\n", - " model_client=model_client,\n", - " tools=[percentage_change_tool],\n", - " system_message=\"\"\"\n", - " You are a data analyst.\n", - " Given the tasks you have been assigned, you should analyze the data and provide results using the tools provided.\n", - " If you have not seen the data, ask for it.\n", - " \"\"\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{note}\n", - "By default, {py:class}`~autogen_agentchat.agents.AssistantAgent` returns the\n", - "tool output as the response. If your tool does not return a well-formed\n", - "string in natural language format, you may want to add a reflection step\n", - "within the agent by setting `reflect_on_tool_use=True` when creating the agent.\n", - "This will allow the agent to reflect on the tool output and provide a natural\n", - "language response.\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Workflow\n", - "\n", - "1. The task is received by the {py:class}`~autogen_agentchat.teams.SelectorGroupChat` which, based on agent descriptions, selects the most appropriate agent to handle the initial task (typically the Planning Agent).\n", - "\n", - "2. The **Planning Agent** analyzes the task and breaks it down into subtasks, assigning each to the most appropriate agent using the format:\n", - " ` : `\n", - "\n", - "3. Based on the conversation context and agent descriptions, the {py:class}`~autogen_agent.teams.SelectorGroupChat` manager dynamically selects the next agent to handle their assigned subtask.\n", - "\n", - "4. The **Web Search Agent** performs searches one at a time, storing results in the shared conversation history.\n", - "\n", - "5. The **Data Analyst** processes the gathered information using available calculation tools when selected.\n", - "\n", - "6. The workflow continues with agents being dynamically selected until either:\n", - " - The Planning Agent determines all subtasks are complete and sends \"TERMINATE\"\n", - " - An alternative termination condition is met (e.g., a maximum number of messages)\n", - "\n", - "When defining your agents, make sure to include a helpful {py:attr}`~autogen_agentchat.base.ChatAgent.description` since this is used to decide which agent to select next." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Termination Conditions\n", - "\n", - "Let's use two termination conditions:\n", - "{py:class}`~autogen_agentchat.conditions.TextMentionTermination` to end the conversation when the Planning Agent sends \"TERMINATE\",\n", - "and {py:class}`~autogen_agentchat.conditions.MaxMessageTermination` to limit the conversation to 25 messages to avoid infinite loop." - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [], - "source": [ - "text_mention_termination = TextMentionTermination(\"TERMINATE\")\n", - "max_messages_termination = MaxMessageTermination(max_messages=25)\n", - "termination = text_mention_termination | max_messages_termination" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Selector Prompt\n", - "\n", - "{py:class}`~autogen_agentchat.teams.SelectorGroupChat` uses a model to select\n", - "the next speaker based on the conversation context.\n", - "We will use a custom selector prompt to properly align with the workflow." - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "metadata": {}, - "outputs": [], - "source": [ - "selector_prompt = \"\"\"Select an agent to perform task.\n", - "\n", - "{roles}\n", - "\n", - "Current conversation context:\n", - "{history}\n", - "\n", - "Read the above conversation, then select an agent from {participants} to perform the next task.\n", - "Make sure the planner agent has assigned tasks before other agents start working.\n", - "Only select one agent.\n", - "\"\"\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{tip}\n", - "Try not to overload the model with too much instruction in the selector prompt.\n", - "\n", - "What is too much? It depends on the capabilities of the model you are using.\n", - "For GPT-4o and equivalents, you can use a selector prompt with a condition for when each speaker should be selected.\n", - "For smaller models such as Phi-4, you should keep the selector prompt as simple as possible\n", - "such as the one used in this example.\n", - "\n", - "Generally, if you find yourself writing multiple conditions for each agent,\n", - "it is a sign that you should consider using a custom selection function,\n", - "or breaking down the task into smaller, sequential tasks to be handled by\n", - "separate agents or teams.\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Running the Team\n", - "\n", - "Let's create the team with the agents, termination conditions, and custom selector prompt." - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "metadata": {}, - "outputs": [], - "source": [ - "team = SelectorGroupChat(\n", - " [planning_agent, web_search_agent, data_analyst_agent],\n", - " model_client=model_client,\n", - " termination_condition=termination,\n", - " selector_prompt=selector_prompt,\n", - " allow_repeated_speaker=True, # Allow an agent to speak multiple turns in a row.\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now we run the team with a task to find information about an NBA player." - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "metadata": {}, - "outputs": [], - "source": [ - "task = \"Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\"" - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n", - "---------- PlanningAgent ----------\n", - "To complete this task, we need to perform the following subtasks:\n", - "\n", - "1. Find out which Miami Heat player had the highest points in the 2006-2007 season.\n", - "2. Gather data on this player's total rebounds for the 2007-2008 season.\n", - "3. Gather data on this player's total rebounds for the 2008-2009 season.\n", - "4. Calculate the percentage change in the player's total rebounds between the 2007-2008 and 2008-2009 seasons.\n", - "\n", - "I'll assign these tasks accordingly:\n", - "\n", - "1. WebSearchAgent: Search for the Miami Heat player with the highest points in the 2006-2007 NBA season.\n", - "2. WebSearchAgent: Find the total rebounds for this player in the 2007-2008 NBA season.\n", - "3. WebSearchAgent: Find the total rebounds for this player in the 2008-2009 NBA season.\n", - "4. DataAnalystAgent: Calculate the percentage change in total rebounds from the 2007-2008 season to the 2008-2009 season for this player.\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_89tUNHaAM0kKQYPJLleGUKK7', arguments='{\"query\":\"Miami Heat player highest points 2006-2007 season\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', name='search_web_tool', call_id='call_89tUNHaAM0kKQYPJLleGUKK7', is_error=False)]\n", - "---------- WebSearchAgent ----------\n", - "Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", - " Udonis Haslem: 844 points\n", - " Dwayne Wade: 1397 points\n", - " James Posey: 550 points\n", - " ...\n", - " \n", - "---------- WebSearchAgent ----------\n", - "The Miami Heat player with the highest points in the 2006-2007 season was Dwyane Wade, with 1,397 points.\n", - "\n", - "Next, I will search for Dwyane Wade's total rebounds for the 2007-2008 season.\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_RC55TkSjG3JXRuVOTPrcE1RL', arguments='{\"query\":\"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', name='search_web_tool', call_id='call_RC55TkSjG3JXRuVOTPrcE1RL', is_error=False)]\n", - "---------- WebSearchAgent ----------\n", - "The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_pBXoABrErDow0rZjw3tjOZol', arguments='{\"query\":\"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', name='search_web_tool', call_id='call_pBXoABrErDow0rZjw3tjOZol', is_error=False)]\n", - "---------- WebSearchAgent ----------\n", - "The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n", - "---------- DataAnalystAgent ----------\n", - "[FunctionCall(id='call_qMxxXtcJsiK8KFSSCx3zm0is', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')]\n", - "---------- DataAnalystAgent ----------\n", - "[FunctionExecutionResult(content='85.98130841121495', name='percentage_change_tool', call_id='call_qMxxXtcJsiK8KFSSCx3zm0is', is_error=False)]\n", - "---------- DataAnalystAgent ----------\n", - "85.98130841121495\n", - "---------- PlanningAgent ----------\n", - "The player with the highest points for the Miami Heat in the 2006-2007 NBA season was Dwyane Wade, who scored 1,397 points. The percentage change in Dwyane Wade's total rebounds from 214 in the 2007-2008 season to 398 in the 2008-2009 season is approximately 85.98%.\n", - "\n", - "TERMINATE\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=161, completion_tokens=220), metadata={}, content=\"To complete this task, we need to perform the following subtasks:\\n\\n1. Find out which Miami Heat player had the highest points in the 2006-2007 season.\\n2. Gather data on this player's total rebounds for the 2007-2008 season.\\n3. Gather data on this player's total rebounds for the 2008-2009 season.\\n4. Calculate the percentage change in the player's total rebounds between the 2007-2008 and 2008-2009 seasons.\\n\\nI'll assign these tasks accordingly:\\n\\n1. WebSearchAgent: Search for the Miami Heat player with the highest points in the 2006-2007 NBA season.\\n2. WebSearchAgent: Find the total rebounds for this player in the 2007-2008 NBA season.\\n3. WebSearchAgent: Find the total rebounds for this player in the 2008-2009 NBA season.\\n4. DataAnalystAgent: Calculate the percentage change in total rebounds from the 2007-2008 season to the 2008-2009 season for this player.\", type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=368, completion_tokens=27), metadata={}, content=[FunctionCall(id='call_89tUNHaAM0kKQYPJLleGUKK7', arguments='{\"query\":\"Miami Heat player highest points 2006-2007 season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', name='search_web_tool', call_id='call_89tUNHaAM0kKQYPJLleGUKK7', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, metadata={}, content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', type='ToolCallSummaryMessage'), ThoughtEvent(source='WebSearchAgent', models_usage=None, metadata={}, content=\"The Miami Heat player with the highest points in the 2006-2007 season was Dwyane Wade, with 1,397 points.\\n\\nNext, I will search for Dwyane Wade's total rebounds for the 2007-2008 season.\", type='ThoughtEvent'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=460, completion_tokens=83), metadata={}, content=[FunctionCall(id='call_RC55TkSjG3JXRuVOTPrcE1RL', arguments='{\"query\":\"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', name='search_web_tool', call_id='call_RC55TkSjG3JXRuVOTPrcE1RL', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, metadata={}, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', type='ToolCallSummaryMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=585, completion_tokens=28), metadata={}, content=[FunctionCall(id='call_pBXoABrErDow0rZjw3tjOZol', arguments='{\"query\":\"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', name='search_web_tool', call_id='call_pBXoABrErDow0rZjw3tjOZol', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, metadata={}, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', type='ToolCallSummaryMessage'), ToolCallRequestEvent(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=496, completion_tokens=21), metadata={}, content=[FunctionCall(id='call_qMxxXtcJsiK8KFSSCx3zm0is', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='DataAnalystAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='85.98130841121495', name='percentage_change_tool', call_id='call_qMxxXtcJsiK8KFSSCx3zm0is', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='DataAnalystAgent', models_usage=None, metadata={}, content='85.98130841121495', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=528, completion_tokens=80), metadata={}, content=\"The player with the highest points for the Miami Heat in the 2006-2007 NBA season was Dwyane Wade, who scored 1,397 points. The percentage change in Dwyane Wade's total rebounds from 214 in the 2007-2008 season to 398 in the 2008-2009 season is approximately 85.98%.\\n\\nTERMINATE\", type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")" - ] - }, - "execution_count": 8, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "# Use asyncio.run(...) if you are running this in a script.\n", - "await Console(team.run_stream(task=task))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As we can see, after the Web Search Agent conducts the necessary searches and the Data Analyst Agent completes the necessary calculations, we find that Dwayne Wade was the Miami Heat player with the highest points in the 2006-2007 season, and the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons is 85.98%!" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Custom Selector Function" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Often times we want better control over the selection process.\n", - "To this end, we can set the `selector_func` argument with a custom selector function to override the default model-based selection.\n", - "This allows us to implement more complex selection logic and state-based transitions.\n", - "\n", - "For instance, we want the Planning Agent to speak immediately after any specialized agent to check the progress.\n", - "\n", - "```{note}\n", - "Returning `None` from the custom selector function will use the default model-based selection.\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n", - "---------- PlanningAgent ----------\n", - "To answer this question, we need to follow these steps: \n", - "\n", - "1. Identify the Miami Heat player with the highest points in the 2006-2007 season.\n", - "2. Retrieve the total rebounds of that player for the 2007-2008 and 2008-2009 seasons.\n", - "3. Calculate the percentage change in his total rebounds between the two seasons.\n", - "\n", - "Let's delegate these tasks:\n", - "\n", - "1. WebSearchAgent: Find the Miami Heat player with the highest points in the 2006-2007 NBA season.\n", - "2. WebSearchAgent: Retrieve the total rebounds for the identified player during the 2007-2008 NBA season.\n", - "3. WebSearchAgent: Retrieve the total rebounds for the identified player during the 2008-2009 NBA season.\n", - "4. DataAnalystAgent: Calculate the percentage change in total rebounds between the 2007-2008 and 2008-2009 seasons for the player found.\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_Pz82ndNLSV4cH0Sg6g7ArP4L', arguments='{\"query\":\"Miami Heat player highest points 2006-2007 season\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', call_id='call_Pz82ndNLSV4cH0Sg6g7ArP4L')]\n", - "---------- WebSearchAgent ----------\n", - "Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", - " Udonis Haslem: 844 points\n", - " Dwayne Wade: 1397 points\n", - " James Posey: 550 points\n", - " ...\n", - " \n", - "---------- PlanningAgent ----------\n", - "Great! Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season. Now, let's continue with the next tasks:\n", - "\n", - "2. WebSearchAgent: Retrieve the total rebounds for Dwyane Wade during the 2007-2008 NBA season.\n", - "3. WebSearchAgent: Retrieve the total rebounds for Dwyane Wade during the 2008-2009 NBA season.\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_3qv9so2DXFZIHtzqDIfXoFID', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool'), FunctionCall(id='call_Vh7zzzWUeiUAvaYjP0If0k1k', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_3qv9so2DXFZIHtzqDIfXoFID'), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_Vh7zzzWUeiUAvaYjP0If0k1k')]\n", - "---------- WebSearchAgent ----------\n", - "The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n", - "The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n", - "---------- PlanningAgent ----------\n", - "Now let's calculate the percentage change in total rebounds between the 2007-2008 and 2008-2009 seasons for Dwyane Wade.\n", - "\n", - "4. DataAnalystAgent: Calculate the percentage change in total rebounds for Dwyane Wade between the 2007-2008 and 2008-2009 seasons.\n", - "---------- DataAnalystAgent ----------\n", - "[FunctionCall(id='call_FXnPSr6JVGfAWs3StIizbt2V', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')]\n", - "---------- DataAnalystAgent ----------\n", - "[FunctionExecutionResult(content='85.98130841121495', call_id='call_FXnPSr6JVGfAWs3StIizbt2V')]\n", - "---------- DataAnalystAgent ----------\n", - "85.98130841121495\n", - "---------- PlanningAgent ----------\n", - "Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring a total of 1397 points. The percentage change in his total rebounds from the 2007-2008 season (214 rebounds) to the 2008-2009 season (398 rebounds) is approximately 86.0%.\n", - "\n", - "TERMINATE\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=161, completion_tokens=192), content=\"To answer this question, we need to follow these steps: \\n\\n1. Identify the Miami Heat player with the highest points in the 2006-2007 season.\\n2. Retrieve the total rebounds of that player for the 2007-2008 and 2008-2009 seasons.\\n3. Calculate the percentage change in his total rebounds between the two seasons.\\n\\nLet's delegate these tasks:\\n\\n1. WebSearchAgent: Find the Miami Heat player with the highest points in the 2006-2007 NBA season.\\n2. WebSearchAgent: Retrieve the total rebounds for the identified player during the 2007-2008 NBA season.\\n3. WebSearchAgent: Retrieve the total rebounds for the identified player during the 2008-2009 NBA season.\\n4. DataAnalystAgent: Calculate the percentage change in total rebounds between the 2007-2008 and 2008-2009 seasons for the player found.\", type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=340, completion_tokens=27), content=[FunctionCall(id='call_Pz82ndNLSV4cH0Sg6g7ArP4L', arguments='{\"query\":\"Miami Heat player highest points 2006-2007 season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', call_id='call_Pz82ndNLSV4cH0Sg6g7ArP4L')], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=420, completion_tokens=87), content=\"Great! Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season. Now, let's continue with the next tasks:\\n\\n2. WebSearchAgent: Retrieve the total rebounds for Dwyane Wade during the 2007-2008 NBA season.\\n3. WebSearchAgent: Retrieve the total rebounds for Dwyane Wade during the 2008-2009 NBA season.\", type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=525, completion_tokens=71), content=[FunctionCall(id='call_3qv9so2DXFZIHtzqDIfXoFID', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool'), FunctionCall(id='call_Vh7zzzWUeiUAvaYjP0If0k1k', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_3qv9so2DXFZIHtzqDIfXoFID'), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_Vh7zzzWUeiUAvaYjP0If0k1k')], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\\nThe number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=569, completion_tokens=68), content=\"Now let's calculate the percentage change in total rebounds between the 2007-2008 and 2008-2009 seasons for Dwyane Wade.\\n\\n4. DataAnalystAgent: Calculate the percentage change in total rebounds for Dwyane Wade between the 2007-2008 and 2008-2009 seasons.\", type='TextMessage'), ToolCallRequestEvent(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=627, completion_tokens=21), content=[FunctionCall(id='call_FXnPSr6JVGfAWs3StIizbt2V', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='DataAnalystAgent', models_usage=None, content=[FunctionExecutionResult(content='85.98130841121495', call_id='call_FXnPSr6JVGfAWs3StIizbt2V')], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='DataAnalystAgent', models_usage=None, content='85.98130841121495', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=659, completion_tokens=76), content='Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring a total of 1397 points. The percentage change in his total rebounds from the 2007-2008 season (214 rebounds) to the 2008-2009 season (398 rebounds) is approximately 86.0%.\\n\\nTERMINATE', type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")" - ] - }, - "execution_count": 10, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "def selector_func(messages: Sequence[AgentEvent | ChatMessage]) -> str | None:\n", - " if messages[-1].source != planning_agent.name:\n", - " return planning_agent.name\n", - " return None\n", - "\n", - "\n", - "# Reset the previous team and run the chat again with the selector function.\n", - "await team.reset()\n", - "team = SelectorGroupChat(\n", - " [planning_agent, web_search_agent, data_analyst_agent],\n", - " model_client=model_client,\n", - " termination_condition=termination,\n", - " selector_prompt=selector_prompt,\n", - " allow_repeated_speaker=True,\n", - " selector_func=selector_func,\n", - ")\n", - "\n", - "await Console(team.run_stream(task=task))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "You can see from the conversation log that the Planning Agent always speaks immediately after the specialized agents.\n", - "\n", - "```{tip}\n", - "Each participant agent only makes one step (executing tools, generating a response, etc.)\n", - "on each turn. \n", - "If you want an {py:class}`~autogen_agentchat.agents.AssistantAgent` to repeat\n", - "until it stop returning a {py:class}`~autogen_agentchat.messages.ToolCallSummaryMessage`\n", - "when it has finished running all the tools it needs to run, you can do so by\n", - "checking the last message and returning the agent if it is a\n", - "{py:class}`~autogen_agentchat.messages.ToolCallSummaryMessage`.\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Custom Candidate Function" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "One more possible requirement might be to automatically select the next speaker from a filtered list of agents.\n", - "For this, we can set `candidate_func` parameter with a custom candidate function to filter down the list of potential agents for speaker selection for each turn of groupchat.\n", - "\n", - "This allow us to restrict speaker selection to a specific set of agents after a given agent.\n", - "\n", - "\n", - "```{note}\n", - "The `candidate_func` is only valid if `selector_func` is not set.\n", - "Returning `None` or an empty list `[]` from the custom candidate function will raise a `ValueError`.\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n", - "---------- PlanningAgent ----------\n", - "To answer this question, we'll break it down into two main subtasks:\n", - "\n", - "1. Identify the Miami Heat player with the highest points in the 2006-2007 season.\n", - "2. Calculate the percentage change in that player's total rebounds between the 2007-2008 and 2008-2009 seasons.\n", - "\n", - "Let's assign these tasks:\n", - "\n", - "1. WebSearchAgent: Search for the Miami Heat player with the highest points in the 2006-2007 NBA season.\n", - "2. WebSearchAgent: Find the total rebound statistics for that identified player for both the 2007-2008 and 2008-2009 NBA seasons.\n", - "3. DataAnalystAgent: Calculate the percentage change in the player's total rebounds between the 2007-2008 and 2008-2009 seasons once the data is retrieved.\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_WtR5KTfEIxs3jIO25gjAw7dF', arguments='{\"query\":\"Miami Heat highest points scorer 2006-2007 NBA season\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', name='search_web_tool', call_id='call_WtR5KTfEIxs3jIO25gjAw7dF', is_error=False)]\n", - "---------- WebSearchAgent ----------\n", - "Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", - " Udonis Haslem: 844 points\n", - " Dwayne Wade: 1397 points\n", - " James Posey: 550 points\n", - " ...\n", - " \n", - "---------- DataAnalystAgent ----------\n", - "[FunctionCall(id='call_9HA3DEacUl4WuG2G2PtRkXAO', arguments='{\"start\": 432, \"end\": 527}', name='percentage_change_tool')]\n", - "---------- DataAnalystAgent ----------\n", - "[FunctionExecutionResult(content='21.99074074074074', name='percentage_change_tool', call_id='call_9HA3DEacUl4WuG2G2PtRkXAO', is_error=False)]\n", - "---------- DataAnalystAgent ----------\n", - "21.99074074074074\n", - "---------- PlanningAgent ----------\n", - "It seems we've missed some context there, so let's assign the subtasks again for clarity:\n", - "\n", - "Based on the search results, Dwyane Wade had the highest points for the Miami Heat in the 2006-2007 season with 1397 points.\n", - "\n", - "Now, let's find the necessary rebound statistics:\n", - "\n", - "2. WebSearchAgent: Find Dwyane Wade's total rebound statistics for both the 2007-2008 and 2008-2009 NBA seasons.\n", - "3. DataAnalystAgent: Once the data is retrieved, calculate the percentage change in Dwyane Wade's total rebounds between the 2007-2008 and 2008-2009 seasons.\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_3i1wTDSjkGg6Ev8YKYWkZK55', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 NBA season\"}', name='search_web_tool'), FunctionCall(id='call_NRAs6jHxXRi8zsvpW5WlHAaU', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 NBA season\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', name='search_web_tool', call_id='call_3i1wTDSjkGg6Ev8YKYWkZK55', is_error=False), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', name='search_web_tool', call_id='call_NRAs6jHxXRi8zsvpW5WlHAaU', is_error=False)]\n", - "---------- WebSearchAgent ----------\n", - "The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n", - "The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n", - "---------- PlanningAgent ----------\n", - "The total rebounds for Dwyane Wade in the 2007-2008 season were 214, and in the 2008-2009 season, they were 398.\n", - "\n", - "Now, let's calculate the percentage change.\n", - "\n", - "3. DataAnalystAgent: Calculate the percentage change in Dwyane Wade's total rebounds from the 2007-2008 season to the 2008-2009 season.\n", - "---------- DataAnalystAgent ----------\n", - "[FunctionCall(id='call_XECA7ezz7VIKbf8IbZYSCSpI', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')]\n", - "---------- DataAnalystAgent ----------\n", - "[FunctionExecutionResult(content='85.98130841121495', name='percentage_change_tool', call_id='call_XECA7ezz7VIKbf8IbZYSCSpI', is_error=False)]\n", - "---------- DataAnalystAgent ----------\n", - "85.98130841121495\n", - "---------- PlanningAgent ----------\n", - "The Miami Heat player with the highest points in the 2006-2007 season was Dwyane Wade, with 1397 points. The percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons was approximately 85.98%.\n", - "\n", - "TERMINATE\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=161, completion_tokens=169), metadata={}, content=\"To answer this question, we'll break it down into two main subtasks:\\n\\n1. Identify the Miami Heat player with the highest points in the 2006-2007 season.\\n2. Calculate the percentage change in that player's total rebounds between the 2007-2008 and 2008-2009 seasons.\\n\\nLet's assign these tasks:\\n\\n1. WebSearchAgent: Search for the Miami Heat player with the highest points in the 2006-2007 NBA season.\\n2. WebSearchAgent: Find the total rebound statistics for that identified player for both the 2007-2008 and 2008-2009 NBA seasons.\\n3. DataAnalystAgent: Calculate the percentage change in the player's total rebounds between the 2007-2008 and 2008-2009 seasons once the data is retrieved.\", type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=324, completion_tokens=28), metadata={}, content=[FunctionCall(id='call_WtR5KTfEIxs3jIO25gjAw7dF', arguments='{\"query\":\"Miami Heat highest points scorer 2006-2007 NBA season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', name='search_web_tool', call_id='call_WtR5KTfEIxs3jIO25gjAw7dF', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, metadata={}, content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', type='ToolCallSummaryMessage'), ToolCallRequestEvent(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=390, completion_tokens=37), metadata={}, content=[FunctionCall(id='call_9HA3DEacUl4WuG2G2PtRkXAO', arguments='{\"start\": 432, \"end\": 527}', name='percentage_change_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='DataAnalystAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='21.99074074074074', name='percentage_change_tool', call_id='call_9HA3DEacUl4WuG2G2PtRkXAO', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='DataAnalystAgent', models_usage=None, metadata={}, content='21.99074074074074', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=413, completion_tokens=137), metadata={}, content=\"It seems we've missed some context there, so let's assign the subtasks again for clarity:\\n\\nBased on the search results, Dwyane Wade had the highest points for the Miami Heat in the 2006-2007 season with 1397 points.\\n\\nNow, let's find the necessary rebound statistics:\\n\\n2. WebSearchAgent: Find Dwyane Wade's total rebound statistics for both the 2007-2008 and 2008-2009 NBA seasons.\\n3. DataAnalystAgent: Once the data is retrieved, calculate the percentage change in Dwyane Wade's total rebounds between the 2007-2008 and 2008-2009 seasons.\", type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=576, completion_tokens=73), metadata={}, content=[FunctionCall(id='call_3i1wTDSjkGg6Ev8YKYWkZK55', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 NBA season\"}', name='search_web_tool'), FunctionCall(id='call_NRAs6jHxXRi8zsvpW5WlHAaU', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 NBA season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', name='search_web_tool', call_id='call_3i1wTDSjkGg6Ev8YKYWkZK55', is_error=False), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', name='search_web_tool', call_id='call_NRAs6jHxXRi8zsvpW5WlHAaU', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, metadata={}, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\\nThe number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=612, completion_tokens=84), metadata={}, content=\"The total rebounds for Dwyane Wade in the 2007-2008 season were 214, and in the 2008-2009 season, they were 398.\\n\\nNow, let's calculate the percentage change.\\n\\n3. DataAnalystAgent: Calculate the percentage change in Dwyane Wade's total rebounds from the 2007-2008 season to the 2008-2009 season.\", type='TextMessage'), ToolCallRequestEvent(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=720, completion_tokens=21), metadata={}, content=[FunctionCall(id='call_XECA7ezz7VIKbf8IbZYSCSpI', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='DataAnalystAgent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='85.98130841121495', name='percentage_change_tool', call_id='call_XECA7ezz7VIKbf8IbZYSCSpI', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='DataAnalystAgent', models_usage=None, metadata={}, content='85.98130841121495', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=718, completion_tokens=63), metadata={}, content='The Miami Heat player with the highest points in the 2006-2007 season was Dwyane Wade, with 1397 points. The percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons was approximately 85.98%.\\n\\nTERMINATE', type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "def candidate_func(messages: Sequence[AgentEvent | ChatMessage]) -> List[str]:\n", - " # keep planning_agent first one to plan out the tasks\n", - " if messages[-1].source == \"user\":\n", - " return [planning_agent.name]\n", - "\n", - " # if previous agent is planning_agent and if it explicitely asks for web_search_agent\n", - " # or data_analyst_agent or both (in-case of re-planning or re-assignment of tasks)\n", - " # then return those specific agents\n", - " last_message = messages[-1]\n", - " if last_message.source == planning_agent.name:\n", - " participants = []\n", - " if web_search_agent.name in last_message.content:\n", - " participants.append(web_search_agent.name)\n", - " if data_analyst_agent.name in last_message.content:\n", - " participants.append(data_analyst_agent.name)\n", - " if participants:\n", - " return participants # SelectorGroupChat will select from the remaining two agents.\n", - "\n", - " # we can assume that the task is finished once the web_search_agent\n", - " # and data_analyst_agent have took their turns, thus we send\n", - " # in planning_agent to terminate the chat\n", - " previous_set_of_agents = set(message.source for message in messages)\n", - " if web_search_agent.name in previous_set_of_agents and data_analyst_agent.name in previous_set_of_agents:\n", - " return [planning_agent.name]\n", - "\n", - " # if no-conditions are met then return all the agents\n", - " return [planning_agent.name, web_search_agent.name, data_analyst_agent.name]\n", - "\n", - "\n", - "# Reset the previous team and run the chat again with the selector function.\n", - "await team.reset()\n", - "team = SelectorGroupChat(\n", - " [planning_agent, web_search_agent, data_analyst_agent],\n", - " model_client=model_client,\n", - " termination_condition=termination,\n", - " candidate_func=candidate_func,\n", - ")\n", - "\n", - "await Console(team.run_stream(task=task))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "You can see from the conversation log that the Planning Agent returns to conversation once the Web Search Agent and Data Analyst Agent took their turns and it finds that the task was not finished as expected so it called the WebSearchAgent again to get rebound values and then called DataAnalysetAgent to get the percentage change." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## User Feedback\n", - "\n", - "We can add {py:class}`~autogen_agentchat.agents.UserProxyAgent` to the team to\n", - "provide user feedback during a run.\n", - "See [Human-in-the-Loop](./tutorial/human-in-the-loop.ipynb) for more details\n", - "about {py:class}`~autogen_agentchat.agents.UserProxyAgent`.\n", - "\n", - "To use the {py:class}`~autogen_agentchat.agents.UserProxyAgent` in the \n", - "web search example, we simply add it to the team and update the selector function\n", - "to always check for user feedback after the planning agent speaks.\n", - "If the user responds with `\"APPROVE\"`, the conversation continues, otherwise,\n", - "the planning agent tries again, until the user approves." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- PlanningAgent ----------\n", - "To address the user's query, we will need to perform the following tasks:\n", - "\n", - "1. Identify the Miami Heat player with the highest points in the 2006-2007 season.\n", - "2. Find the total rebounds for that player in the 2007-2008 season.\n", - "3. Find the total rebounds for that player in the 2008-2009 season.\n", - "4. Calculate the percentage change in the total rebounds between the 2007-2008 and 2008-2009 seasons.\n", - "\n", - "Let's assign these tasks:\n", - "\n", - "1. **WebSearchAgent**: Identify the Miami Heat player with the highest points in the 2006-2007 season.\n", - " \n", - "(Task 2 and 3 depend on the result of Task 1. We'll proceed with Tasks 2 and 3 once Task 1 is complete.)\n", - "---------- UserProxyAgent ----------\n", - "approve\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_0prr3fUnG5CtisUG7QeygW0w', arguments='{\"query\":\"Miami Heat highest points scorer 2006-2007 NBA season\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', call_id='call_0prr3fUnG5CtisUG7QeygW0w')]\n", - "---------- WebSearchAgent ----------\n", - "Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", - " Udonis Haslem: 844 points\n", - " Dwayne Wade: 1397 points\n", - " James Posey: 550 points\n", - " ...\n", - " \n", - "---------- PlanningAgent ----------\n", - "Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring 1397 points.\n", - "\n", - "Next, we need to find Dwyane Wade's total rebounds for the 2007-2008 and 2008-2009 seasons:\n", - "\n", - "2. **WebSearchAgent**: Find Dwyane Wade's total rebounds for the 2007-2008 season.\n", - "3. **WebSearchAgent**: Find Dwyane Wade's total rebounds for the 2008-2009 season.\n", - "---------- UserProxyAgent ----------\n", - "approve\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_fBZe80NaBfruOVGwRWbhXyRm', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 NBA season\"}', name='search_web_tool'), FunctionCall(id='call_cURYibna4fGxySiL7IYt0c3s', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 NBA season\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_fBZe80NaBfruOVGwRWbhXyRm'), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_cURYibna4fGxySiL7IYt0c3s')]\n", - "---------- WebSearchAgent ----------\n", - "The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n", - "The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n", - "---------- PlanningAgent ----------\n", - "Now that we have Dwyane Wade's total rebounds for both seasons, we can calculate the percentage change:\n", - "\n", - "4. **DataAnalystAgent**: Calculate the percentage change in Dwyane Wade's total rebounds from the 2007-2008 season (214 rebounds) to the 2008-2009 season (398 rebounds).\n", - "---------- UserProxyAgent ----------\n", - "approve\n", - "---------- DataAnalystAgent ----------\n", - "[FunctionCall(id='call_z3uog7t2x0z1Suzl5hACF9hY', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')]\n", - "---------- DataAnalystAgent ----------\n", - "[FunctionExecutionResult(content='85.98130841121495', call_id='call_z3uog7t2x0z1Suzl5hACF9hY')]\n", - "---------- DataAnalystAgent ----------\n", - "85.98130841121495\n", - "---------- PlanningAgent ----------\n", - "Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring 1397 points. His total rebounds increased from 214 in the 2007-2008 season to 398 in the 2008-2009 season, which is a percentage change of approximately 85.98%.\n", - "\n", - "TERMINATE\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=161, completion_tokens=166), content=\"To address the user's query, we will need to perform the following tasks:\\n\\n1. Identify the Miami Heat player with the highest points in the 2006-2007 season.\\n2. Find the total rebounds for that player in the 2007-2008 season.\\n3. Find the total rebounds for that player in the 2008-2009 season.\\n4. Calculate the percentage change in the total rebounds between the 2007-2008 and 2008-2009 seasons.\\n\\nLet's assign these tasks:\\n\\n1. **WebSearchAgent**: Identify the Miami Heat player with the highest points in the 2006-2007 season.\\n \\n(Task 2 and 3 depend on the result of Task 1. We'll proceed with Tasks 2 and 3 once Task 1 is complete.)\", type='TextMessage'), UserInputRequestedEvent(source='UserProxyAgent', models_usage=None, request_id='2a433f88-f886-4b39-a078-ea1acdcb2f9d', content='', type='UserInputRequestedEvent'), TextMessage(source='UserProxyAgent', models_usage=None, content='approve', type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=323, completion_tokens=28), content=[FunctionCall(id='call_0prr3fUnG5CtisUG7QeygW0w', arguments='{\"query\":\"Miami Heat highest points scorer 2006-2007 NBA season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', call_id='call_0prr3fUnG5CtisUG7QeygW0w')], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=403, completion_tokens=112), content=\"Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring 1397 points.\\n\\nNext, we need to find Dwyane Wade's total rebounds for the 2007-2008 and 2008-2009 seasons:\\n\\n2. **WebSearchAgent**: Find Dwyane Wade's total rebounds for the 2007-2008 season.\\n3. **WebSearchAgent**: Find Dwyane Wade's total rebounds for the 2008-2009 season.\", type='TextMessage'), UserInputRequestedEvent(source='UserProxyAgent', models_usage=None, request_id='23dd4570-2391-41e9-aeea-86598499792c', content='', type='UserInputRequestedEvent'), TextMessage(source='UserProxyAgent', models_usage=None, content='approve', type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=543, completion_tokens=73), content=[FunctionCall(id='call_fBZe80NaBfruOVGwRWbhXyRm', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 NBA season\"}', name='search_web_tool'), FunctionCall(id='call_cURYibna4fGxySiL7IYt0c3s', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 NBA season\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_fBZe80NaBfruOVGwRWbhXyRm'), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_cURYibna4fGxySiL7IYt0c3s')], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\\nThe number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=586, completion_tokens=70), content=\"Now that we have Dwyane Wade's total rebounds for both seasons, we can calculate the percentage change:\\n\\n4. **DataAnalystAgent**: Calculate the percentage change in Dwyane Wade's total rebounds from the 2007-2008 season (214 rebounds) to the 2008-2009 season (398 rebounds).\", type='TextMessage'), UserInputRequestedEvent(source='UserProxyAgent', models_usage=None, request_id='e849d193-4ab3-4558-8560-7dbc062a0aee', content='', type='UserInputRequestedEvent'), TextMessage(source='UserProxyAgent', models_usage=None, content='approve', type='TextMessage'), ToolCallRequestEvent(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=655, completion_tokens=21), content=[FunctionCall(id='call_z3uog7t2x0z1Suzl5hACF9hY', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='DataAnalystAgent', models_usage=None, content=[FunctionExecutionResult(content='85.98130841121495', call_id='call_z3uog7t2x0z1Suzl5hACF9hY')], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='DataAnalystAgent', models_usage=None, content='85.98130841121495', type='ToolCallSummaryMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=687, completion_tokens=74), content='Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring 1397 points. His total rebounds increased from 214 in the 2007-2008 season to 398 in the 2008-2009 season, which is a percentage change of approximately 85.98%.\\n\\nTERMINATE', type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")" - ] - }, - "execution_count": 9, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "user_proxy_agent = UserProxyAgent(\"UserProxyAgent\", description=\"A proxy for the user to approve or disapprove tasks.\")\n", - "\n", - "\n", - "def selector_func_with_user_proxy(messages: Sequence[AgentEvent | ChatMessage]) -> str | None:\n", - " if messages[-1].source != planning_agent.name and messages[-1].source != user_proxy_agent.name:\n", - " # Planning agent should be the first to engage when given a new task, or check progress.\n", - " return planning_agent.name\n", - " if messages[-1].source == planning_agent.name:\n", - " if messages[-2].source == user_proxy_agent.name and \"APPROVE\" in messages[-1].content.upper(): # type: ignore\n", - " # User has approved the plan, proceed to the next agent.\n", - " return None\n", - " # Use the user proxy agent to get the user's approval to proceed.\n", - " return user_proxy_agent.name\n", - " if messages[-1].source == user_proxy_agent.name:\n", - " # If the user does not approve, return to the planning agent.\n", - " if \"APPROVE\" not in messages[-1].content.upper(): # type: ignore\n", - " return planning_agent.name\n", - " return None\n", - "\n", - "\n", - "# Reset the previous agents and run the chat again with the user proxy agent and selector function.\n", - "await team.reset()\n", - "team = SelectorGroupChat(\n", - " [planning_agent, web_search_agent, data_analyst_agent, user_proxy_agent],\n", - " model_client=model_client,\n", - " termination_condition=termination,\n", - " selector_prompt=selector_prompt,\n", - " selector_func=selector_func_with_user_proxy,\n", - " allow_repeated_speaker=True,\n", - ")\n", - "\n", - "await Console(team.run_stream(task=task))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now, the user's feedback is incorporated into the conversation flow,\n", - "and the user can approve or reject the planning agent's decisions." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Using Reasoning Models\n", - "\n", - "So far in the examples, we have used a `gpt-4o` model. Models like `gpt-4o`\n", - "and `gemini-1.5-flash` are great at following instructions, so you can\n", - "have relatively detailed instructions in the selector prompt for the team and the \n", - "system messages for each agent to guide their behavior.\n", - "\n", - "However, if you are using a reasoning model like `o3-mini`, you will need to\n", - "keep the selector prompt and system messages as simple and to the point as possible.\n", - "This is because the reasoning models are already good at coming up with their own \n", - "instructions given the context provided to them.\n", - "\n", - "This also means that we don't need a planning agent to break down the task\n", - "anymore, since the {py:class}`~autogen_agentchat.teams.SelectorGroupChat` that\n", - "uses a reasoning model can do that on its own.\n", - "\n", - "In the following example, we will use `o3-mini` as the model for the\n", - "agents and the team, and we will not use a planning agent.\n", - "Also, we are keeping the selector prompt and system messages as simple as possible." - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "metadata": {}, - "outputs": [], - "source": [ - "model_client = OpenAIChatCompletionClient(model=\"o3-mini\")\n", - "\n", - "web_search_agent = AssistantAgent(\n", - " \"WebSearchAgent\",\n", - " description=\"An agent for searching information on the web.\",\n", - " tools=[search_web_tool],\n", - " model_client=model_client,\n", - " system_message=\"\"\"Use web search tool to find information.\"\"\",\n", - ")\n", - "\n", - "data_analyst_agent = AssistantAgent(\n", - " \"DataAnalystAgent\",\n", - " description=\"An agent for performing calculations.\",\n", - " model_client=model_client,\n", - " tools=[percentage_change_tool],\n", - " system_message=\"\"\"Use tool to perform calculation. If you have not seen the data, ask for it.\"\"\",\n", - ")\n", - "\n", - "user_proxy_agent = UserProxyAgent(\n", - " \"UserProxyAgent\",\n", - " description=\"A user to approve or disapprove tasks.\",\n", - ")\n", - "\n", - "selector_prompt = \"\"\"Select an agent to perform task.\n", - "\n", - "{roles}\n", - "\n", - "Current conversation context:\n", - "{history}\n", - "\n", - "Read the above conversation, then select an agent from {participants} to perform the next task.\n", - "When the task is complete, let the user approve or disapprove the task.\n", - "\"\"\"\n", - "\n", - "team = SelectorGroupChat(\n", - " [web_search_agent, data_analyst_agent, user_proxy_agent],\n", - " model_client=model_client,\n", - " termination_condition=termination, # Use the same termination condition as before.\n", - " selector_prompt=selector_prompt,\n", - " allow_repeated_speaker=True,\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": 9, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_hl7EP6Lp5jj5wEdxeNHTwUVG', arguments='{\"query\": \"Who was the Miami Heat player with the highest points in the 2006-2007 season Miami Heat statistics Dwyane Wade rebounds percentage change 2007-2008 2008-2009 seasons\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', call_id='call_hl7EP6Lp5jj5wEdxeNHTwUVG', is_error=False)]\n", - "---------- WebSearchAgent ----------\n", - "Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", - " Udonis Haslem: 844 points\n", - " Dwayne Wade: 1397 points\n", - " James Posey: 550 points\n", - " ...\n", - " \n", - "---------- DataAnalystAgent ----------\n", - "I found that in the 2006–2007 season the player with the highest points was Dwyane Wade (with 1,397 points). Could you please provide Dwyane Wade’s total rebounds for the 2007–2008 and the 2008–2009 seasons so I can calculate the percentage change?\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_lppGTILXDvO9waPwKO66ehK6', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 and 2008-2009 seasons for Miami Heat\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_lppGTILXDvO9waPwKO66ehK6', is_error=False)]\n", - "---------- WebSearchAgent ----------\n", - "The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n", - "---------- DataAnalystAgent ----------\n", - "Could you please provide Dwyane Wade’s total rebounds in the 2008-2009 season?\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_r8DBcbJtQfdtugLtyTrqOvoK', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season Miami Heat\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_r8DBcbJtQfdtugLtyTrqOvoK', is_error=False)]\n", - "---------- WebSearchAgent ----------\n", - "The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n", - "---------- DataAnalystAgent ----------\n", - "[FunctionCall(id='call_4jejv1wM7V1osbBCxJze8aQM', arguments='{\"start\": 214, \"end\": 398}', name='percentage_change_tool')]\n", - "---------- DataAnalystAgent ----------\n", - "[FunctionExecutionResult(content='85.98130841121495', call_id='call_4jejv1wM7V1osbBCxJze8aQM', is_error=False)]\n", - "---------- DataAnalystAgent ----------\n", - "85.98130841121495\n", - "---------- DataAnalystAgent ----------\n", - "Dwyane Wade was the Miami Heat player with the highest total points (1,397) during the 2006-2007 season. His total rebounds increased by approximately 86% from 214 in the 2007-2008 season to 398 in the 2008-2009 season.\n", - "---------- UserProxyAgent ----------\n", - "Approve. TERMINATE\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?', type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=103, completion_tokens=384), content=[FunctionCall(id='call_hl7EP6Lp5jj5wEdxeNHTwUVG', arguments='{\"query\": \"Who was the Miami Heat player with the highest points in the 2006-2007 season Miami Heat statistics Dwyane Wade rebounds percentage change 2007-2008 2008-2009 seasons\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', call_id='call_hl7EP6Lp5jj5wEdxeNHTwUVG', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', type='ToolCallSummaryMessage'), TextMessage(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=183, completion_tokens=1038), content='I found that in the 2006–2007 season the player with the highest points was Dwyane Wade (with 1,397 points). Could you please provide Dwyane Wade’s total rebounds for the 2007–2008 and the 2008–2009 seasons so I can calculate the percentage change?', type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=299, completion_tokens=109), content=[FunctionCall(id='call_lppGTILXDvO9waPwKO66ehK6', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 and 2008-2009 seasons for Miami Heat\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_lppGTILXDvO9waPwKO66ehK6', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', type='ToolCallSummaryMessage'), TextMessage(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=291, completion_tokens=224), content='Could you please provide Dwyane Wade’s total rebounds in the 2008-2009 season?', type='TextMessage'), ToolCallRequestEvent(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=401, completion_tokens=37), content=[FunctionCall(id='call_r8DBcbJtQfdtugLtyTrqOvoK', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season Miami Heat\"}', name='search_web_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_r8DBcbJtQfdtugLtyTrqOvoK', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='WebSearchAgent', models_usage=None, content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', type='ToolCallSummaryMessage'), ToolCallRequestEvent(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=353, completion_tokens=158), content=[FunctionCall(id='call_4jejv1wM7V1osbBCxJze8aQM', arguments='{\"start\": 214, \"end\": 398}', name='percentage_change_tool')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='DataAnalystAgent', models_usage=None, content=[FunctionExecutionResult(content='85.98130841121495', call_id='call_4jejv1wM7V1osbBCxJze8aQM', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='DataAnalystAgent', models_usage=None, content='85.98130841121495', type='ToolCallSummaryMessage'), TextMessage(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=394, completion_tokens=138), content='Dwyane Wade was the Miami Heat player with the highest total points (1,397) during the 2006-2007 season. His total rebounds increased by approximately 86% from 214 in the 2007-2008 season to 398 in the 2008-2009 season.', type='TextMessage'), UserInputRequestedEvent(source='UserProxyAgent', models_usage=None, request_id='b3b05408-73fc-47d4-b832-16c9f447cd6e', content='', type='UserInputRequestedEvent'), TextMessage(source='UserProxyAgent', models_usage=None, content='Approve. TERMINATE', type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")" - ] - }, - "execution_count": 9, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "await Console(team.run_stream(task=task))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{tip}\n", - "For more guidance on how to prompt reasoning models, see the\n", - "Azure AI Services Blog on [Prompt Engineering for OpenAI's O1 and O3-mini Reasoning Models](https://techcommunity.microsoft.com/blog/azure-ai-services-blog/prompt-engineering-for-openai%E2%80%99s-o1-and-o3-mini-reasoning-models/4374010)\n", - "```" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.3" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tracing.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tracing.ipynb deleted file mode 100644 index 8e6f07b90482..000000000000 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tracing.ipynb +++ /dev/null @@ -1,403 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Tracing and Observability\n", - "\n", - "AutoGen has [built-in support for tracing](https://microsoft.github.io/autogen/dev/user-guide/core-user-guide/framework/telemetry.html) and observability for collecting comprehensive records on the execution of your application. This feature is useful for debugging, performance analysis, and understanding the flow of your application.\n", - "\n", - "This capability is powered by the [OpenTelemetry](https://opentelemetry.io/) library, which means you can use any OpenTelemetry-compatible backend to collect and analyze traces.\n", - "\n", - "## Setup\n", - "\n", - "To begin, you need to install the OpenTelemetry Python package. You can do this using pip:\n", - "\n", - "```bash\n", - "pip install opentelemetry-sdk\n", - "```\n", - "\n", - "Once you have the SDK installed, the simplest way to set up tracing in AutoGen is to:\n", - "\n", - "1. Configure an OpenTelemetry tracer provider\n", - "2. Set up an exporter to send traces to your backend\n", - "3. Connect the tracer provider to the AutoGen runtime\n", - "\n", - "## Telemetry Backend\n", - "\n", - "To collect and view traces, you need to set up a telemetry backend. Several open-source options are available, including Jaeger, Zipkin. For this example, we will use Jaeger as our telemetry backend.\n", - "\n", - "For a quick start, you can run Jaeger locally using Docker:\n", - "\n", - "```bash\n", - "docker run -d --name jaeger \\\n", - " -e COLLECTOR_OTLP_ENABLED=true \\\n", - " -p 16686:16686 \\\n", - " -p 4317:4317 \\\n", - " -p 4318:4318 \\\n", - " jaegertracing/all-in-one:latest\n", - "```\n", - "\n", - "This command starts a Jaeger instance that listens on port 16686 for the Jaeger UI and port 4317 for the OpenTelemetry collector. You can access the Jaeger UI at `http://localhost:16686`.\n", - "\n", - "## Instrumenting an AgentChat Team\n", - "\n", - "In the following section, we will review how to enable tracing with an AutoGen GroupChat team. The AutoGen runtime already supports open telemetry (automatically logging message metadata). To begin, we will create a tracing service that will be used to instrument the AutoGen runtime. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from opentelemetry import trace\n", - "from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter\n", - "from opentelemetry.sdk.resources import Resource\n", - "from opentelemetry.sdk.trace import TracerProvider\n", - "from opentelemetry.sdk.trace.export import BatchSpanProcessor\n", - "\n", - "otel_exporter = OTLPSpanExporter(endpoint=\"http://localhost:4317\", insecure=True)\n", - "tracer_provider = TracerProvider(resource=Resource({\"service.name\": \"autogen-test-agentchat\"}))\n", - "span_processor = BatchSpanProcessor(otel_exporter)\n", - "tracer_provider.add_span_processor(span_processor)\n", - "trace.set_tracer_provider(tracer_provider)\n", - "\n", - "# we will get reference this tracer later using its service name\n", - "# tracer = trace.get_tracer(\"autogen-test-agentchat\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "\n", - "\n", - "All of the code to create a [team](./tutorial/teams.ipynb) should already be familiar to you. An important note here is that all AgentChat agents and teams are run using the AutoGen core API runtime. In turn, the runtime is already instrumented to log [runtime messaging events (metadata)] (https://github.com/microsoft/autogen/blob/main/python/packages/autogen-core/src/autogen_core/_telemetry/_tracing_config.py) including:\n", - "\n", - "- **create**: When a message is created\n", - "- **send**: When a message is sent\n", - "- **publish**: When a message is published\n", - "- **receive**: When a message is received\n", - "- **intercept**: When a message is intercepted\n", - "- **process**: When a message is processed\n", - "- **ack**: When a message is acknowledged \n", - " " - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": {}, - "outputs": [], - "source": [ - "from autogen_agentchat.agents import AssistantAgent\n", - "from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination\n", - "from autogen_agentchat.teams import SelectorGroupChat\n", - "from autogen_agentchat.ui import Console\n", - "from autogen_core import SingleThreadedAgentRuntime\n", - "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", - "\n", - "\n", - "def search_web_tool(query: str) -> str:\n", - " if \"2006-2007\" in query:\n", - " return \"\"\"Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", - " Udonis Haslem: 844 points\n", - " Dwayne Wade: 1397 points\n", - " James Posey: 550 points\n", - " ...\n", - " \"\"\"\n", - " elif \"2007-2008\" in query:\n", - " return \"The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\"\n", - " elif \"2008-2009\" in query:\n", - " return \"The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\"\n", - " return \"No data found.\"\n", - "\n", - "\n", - "def percentage_change_tool(start: float, end: float) -> float:\n", - " return ((end - start) / start) * 100\n", - "\n", - "\n", - "async def main() -> None:\n", - " model_client = OpenAIChatCompletionClient(model=\"gpt-4o\")\n", - "\n", - " planning_agent = AssistantAgent(\n", - " \"PlanningAgent\",\n", - " description=\"An agent for planning tasks, this agent should be the first to engage when given a new task.\",\n", - " model_client=model_client,\n", - " system_message=\"\"\"\n", - " You are a planning agent.\n", - " Your job is to break down complex tasks into smaller, manageable subtasks.\n", - " Your team members are:\n", - " WebSearchAgent: Searches for information\n", - " DataAnalystAgent: Performs calculations\n", - "\n", - " You only plan and delegate tasks - you do not execute them yourself.\n", - "\n", - " When assigning tasks, use this format:\n", - " 1. : \n", - "\n", - " After all tasks are complete, summarize the findings and end with \"TERMINATE\".\n", - " \"\"\",\n", - " )\n", - "\n", - " web_search_agent = AssistantAgent(\n", - " \"WebSearchAgent\",\n", - " description=\"An agent for searching information on the web.\",\n", - " tools=[search_web_tool],\n", - " model_client=model_client,\n", - " system_message=\"\"\"\n", - " You are a web search agent.\n", - " Your only tool is search_tool - use it to find information.\n", - " You make only one search call at a time.\n", - " Once you have the results, you never do calculations based on them.\n", - " \"\"\",\n", - " )\n", - "\n", - " data_analyst_agent = AssistantAgent(\n", - " \"DataAnalystAgent\",\n", - " description=\"An agent for performing calculations.\",\n", - " model_client=model_client,\n", - " tools=[percentage_change_tool],\n", - " system_message=\"\"\"\n", - " You are a data analyst.\n", - " Given the tasks you have been assigned, you should analyze the data and provide results using the tools provided.\n", - " If you have not seen the data, ask for it.\n", - " \"\"\",\n", - " )\n", - "\n", - " text_mention_termination = TextMentionTermination(\"TERMINATE\")\n", - " max_messages_termination = MaxMessageTermination(max_messages=25)\n", - " termination = text_mention_termination | max_messages_termination\n", - "\n", - " selector_prompt = \"\"\"Select an agent to perform task.\n", - "\n", - " {roles}\n", - "\n", - " Current conversation context:\n", - " {history}\n", - "\n", - " Read the above conversation, then select an agent from {participants} to perform the next task.\n", - " Make sure the planner agent has assigned tasks before other agents start working.\n", - " Only select one agent.\n", - " \"\"\"\n", - "\n", - " task = \"Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\"\n", - "\n", - " tracer = trace.get_tracer(\"autogen-test-agentchat\")\n", - " with tracer.start_as_current_span(\"runtime\"):\n", - " team = SelectorGroupChat(\n", - " [planning_agent, web_search_agent, data_analyst_agent],\n", - " model_client=model_client,\n", - " termination_condition=termination,\n", - " selector_prompt=selector_prompt,\n", - " allow_repeated_speaker=True,\n", - " )\n", - " await Console(team.run_stream(task=task))\n", - "\n", - " await model_client.close()\n", - "\n", - "\n", - "# asyncio.run(main())" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n", - "---------- PlanningAgent ----------\n", - "To accomplish this, we can break down the tasks as follows:\n", - "\n", - "1. WebSearchAgent: Search for the Miami Heat player with the highest points during the 2006-2007 NBA season.\n", - "2. WebSearchAgent: Find the total rebounds for the identified player in both the 2007-2008 and 2008-2009 NBA seasons.\n", - "3. DataAnalystAgent: Calculate the percentage change in total rebounds for the player between the 2007-2008 and 2008-2009 seasons.\n", - "\n", - "Once these tasks are complete, I will summarize the findings.\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_PUhxZyR0CTlWCY4uwd5Zh3WO', arguments='{\"query\":\"Miami Heat highest points scorer 2006-2007 season\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n Udonis Haslem: 844 points\\n Dwayne Wade: 1397 points\\n James Posey: 550 points\\n ...\\n ', name='search_web_tool', call_id='call_PUhxZyR0CTlWCY4uwd5Zh3WO', is_error=False)]\n", - "---------- WebSearchAgent ----------\n", - "Here are the total points scored by Miami Heat players in the 2006-2007 season:\n", - " Udonis Haslem: 844 points\n", - " Dwayne Wade: 1397 points\n", - " James Posey: 550 points\n", - " ...\n", - " \n", - "---------- WebSearchAgent ----------\n", - "Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring 1,397 points. Now, let's find his total rebounds for the 2007-2008 and 2008-2009 NBA seasons.\n", - "---------- WebSearchAgent ----------\n", - "[FunctionCall(id='call_GL7KkWKj9ejIM8FfpgXe2dPk', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool'), FunctionCall(id='call_X81huZoiA30zIjSAIDgb8ebe', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')]\n", - "---------- WebSearchAgent ----------\n", - "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', name='search_web_tool', call_id='call_GL7KkWKj9ejIM8FfpgXe2dPk', is_error=False), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', name='search_web_tool', call_id='call_X81huZoiA30zIjSAIDgb8ebe', is_error=False)]\n", - "---------- WebSearchAgent ----------\n", - "The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n", - "The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n", - "---------- DataAnalystAgent ----------\n", - "[FunctionCall(id='call_kB50RkFVqHptA7FOf0lL2RS8', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')]\n", - "---------- DataAnalystAgent ----------\n", - "[FunctionExecutionResult(content='85.98130841121495', name='percentage_change_tool', call_id='call_kB50RkFVqHptA7FOf0lL2RS8', is_error=False)]\n", - "---------- DataAnalystAgent ----------\n", - "85.98130841121495\n", - "---------- PlanningAgent ----------\n", - "The Miami Heat player with the highest points during the 2006-2007 NBA season was Dwayne Wade, who scored 1,397 points. The percentage increase in his total rebounds from the 2007-2008 season (214 rebounds) to the 2008-2009 season (398 rebounds) was approximately 86%.\n", - "\n", - "TERMINATE\n" - ] - } - ], - "source": [ - "await main()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "You can then use the Jaeger UI to view the traces collected from the application run above. \n", - "\n", - "![Jaeger UI](jaeger.png)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Custom Traces \n", - "\n", - "So far, we are logging only the default events that are generated by the AutoGen runtime (message created, publish etc). However, you can also create custom spans to log specific events in your application. \n", - "\n", - "In the example below, we will show how to log messages from the `RoundRobinGroupChat` team as they are generated by adding custom spans around the team to log runtime events and spans to log messages generated by the team.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "-- primary_agent -- : Leaves cascade like gold, \n", - "Whispering winds cool the earth.\n", - "primary_agent: Leaves cascade like gold, \n", - "Whispering winds cool the earth.\n", - "\n", - "-- critic_agent -- : Your haiku beautifully captures the essence of the fall season with vivid imagery. However, it appears to have six syllables in the second line, which should traditionally be five. Here's a revised version keeping the 5-7-5 syllable structure:\n", - "\n", - "Leaves cascade like gold, \n", - "Whispering winds cool the air. \n", - "\n", - "Please adjust the second line to reflect a five-syllable count. Thank you!\n", - "critic_agent: Your haiku beautifully captures the essence of the fall season with vivid imagery. However, it appears to have six syllables in the second line, which should traditionally be five. Here's a revised version keeping the 5-7-5 syllable structure:\n", - "\n", - "Leaves cascade like gold, \n", - "Whispering winds cool the air. \n", - "\n", - "Please adjust the second line to reflect a five-syllable count. Thank you!\n", - "\n", - "-- primary_agent -- : Leaves cascade like gold, \n", - "Whispering winds cool the air.\n", - "primary_agent: Leaves cascade like gold, \n", - "Whispering winds cool the air.\n", - "\n", - "-- critic_agent -- : APPROVE\n", - "critic_agent: APPROVE\n" - ] - } - ], - "source": [ - "from autogen_agentchat.base import TaskResult\n", - "from autogen_agentchat.conditions import ExternalTermination\n", - "from autogen_agentchat.teams import RoundRobinGroupChat\n", - "from autogen_core import CancellationToken\n", - "\n", - "\n", - "async def run_agents() -> None:\n", - " # Create an OpenAI model client.\n", - " model_client = OpenAIChatCompletionClient(model=\"gpt-4o-2024-08-06\")\n", - "\n", - " # Create the primary agent.\n", - " primary_agent = AssistantAgent(\n", - " \"primary_agent\",\n", - " model_client=model_client,\n", - " system_message=\"You are a helpful AI assistant.\",\n", - " )\n", - "\n", - " # Create the critic agent.\n", - " critic_agent = AssistantAgent(\n", - " \"critic_agent\",\n", - " model_client=model_client,\n", - " system_message=\"Provide constructive feedback. Respond with 'APPROVE' to when your feedbacks are addressed.\",\n", - " )\n", - "\n", - " # Define a termination condition that stops the task if the critic approves.\n", - " text_termination = TextMentionTermination(\"APPROVE\")\n", - "\n", - " tracer = trace.get_tracer(\"autogen-test-agentchat\")\n", - " with tracer.start_as_current_span(\"runtime_round_robin_events\"):\n", - " team = RoundRobinGroupChat([primary_agent, critic_agent], termination_condition=text_termination)\n", - "\n", - " response_stream = team.run_stream(task=\"Write a 2 line haiku about the fall season\")\n", - " async for response in response_stream:\n", - " async for response in response_stream:\n", - " if not isinstance(response, TaskResult):\n", - " print(f\"\\n-- {response.source} -- : {response.content}\")\n", - " with tracer.start_as_current_span(f\"agent_message.{response.source}\") as message_span:\n", - " content = response.content if isinstance(response.content, str) else str(response.content)\n", - " message_span.set_attribute(\"agent.name\", response.source)\n", - " message_span.set_attribute(\"message.content\", content)\n", - " print(f\"{response.source}: {response.content}\")\n", - "\n", - " await model_client.close()\n", - "\n", - "\n", - "await run_agents()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "\n", - "In the code above, we create a new span for each message sent by the agent. We set attributes on the span to include the agent's name and the message content. This allows us to trace the flow of messages through our application and understand how they are processed." - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.9" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/agents.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/agents.ipynb deleted file mode 100644 index 35c9052dee75..000000000000 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/agents.ipynb +++ /dev/null @@ -1,847 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Agents\n", - "\n", - "AutoGen AgentChat provides a set of preset Agents, each with variations in how an agent might respond to messages.\n", - "All agents share the following attributes and methods:\n", - "\n", - "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.name`: The unique name of the agent.\n", - "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.description`: The description of the agent in text.\n", - "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages`: Send the agent a sequence of {py:class}`~autogen_agentchat.messages.ChatMessage` get a {py:class}`~autogen_agentchat.base.Response`. **It is important to note that agents are expected to be stateful and this method is expected to be called with new messages, not the complete history**.\n", - "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream`: Same as {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages` but returns an iterator of {py:class}`~autogen_agentchat.messages.AgentEvent` or {py:class}`~autogen_agentchat.messages.ChatMessage` followed by a {py:class}`~autogen_agentchat.base.Response` as the last item.\n", - "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_reset`: Reset the agent to its initial state.\n", - "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run` and {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run_stream`: convenience methods that call {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages` and {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream` respectively but offer the same interface as [Teams](./teams.ipynb).\n", - "\n", - "See {py:mod}`autogen_agentchat.messages` for more information on AgentChat message types.\n", - "\n", - "\n", - "## Assistant Agent\n", - "\n", - "{py:class}`~autogen_agentchat.agents.AssistantAgent` is a built-in agent that\n", - "uses a language model and has the ability to use tools." - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [], - "source": [ - "from autogen_agentchat.agents import AssistantAgent\n", - "from autogen_agentchat.messages import TextMessage\n", - "from autogen_agentchat.ui import Console\n", - "from autogen_core import CancellationToken\n", - "from autogen_ext.models.openai import OpenAIChatCompletionClient" - ] - }, - { - "cell_type": "code", - "execution_count": 13, - "metadata": {}, - "outputs": [], - "source": [ - "# Define a tool that searches the web for information.\n", - "async def web_search(query: str) -> str:\n", - " \"\"\"Find information on the web\"\"\"\n", - " return \"AutoGen is a programming framework for building multi-agent applications.\"\n", - "\n", - "\n", - "# Create an agent that uses the OpenAI GPT-4o model.\n", - "model_client = OpenAIChatCompletionClient(\n", - " model=\"gpt-4o\",\n", - " # api_key=\"YOUR_API_KEY\",\n", - ")\n", - "agent = AssistantAgent(\n", - " name=\"assistant\",\n", - " model_client=model_client,\n", - " tools=[web_search],\n", - " system_message=\"Use tools to solve tasks.\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "\n", - "## Getting Responses\n", - "\n", - "We can use the {py:meth}`~autogen_agentchat.agents.AssistantAgent.on_messages` method to get the agent response to a given message.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[ToolCallRequestEvent(source='assistant', models_usage=RequestUsage(prompt_tokens=598, completion_tokens=16), content=[FunctionCall(id='call_9UWYM1CgE3ZbnJcSJavNDB79', arguments='{\"query\":\"AutoGen\"}', name='web_search')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='assistant', models_usage=None, content=[FunctionExecutionResult(content='AutoGen is a programming framework for building multi-agent applications.', call_id='call_9UWYM1CgE3ZbnJcSJavNDB79', is_error=False)], type='ToolCallExecutionEvent')]\n", - "source='assistant' models_usage=None content='AutoGen is a programming framework for building multi-agent applications.' type='ToolCallSummaryMessage'\n" - ] - } - ], - "source": [ - "async def assistant_run() -> None:\n", - " response = await agent.on_messages(\n", - " [TextMessage(content=\"Find information on AutoGen\", source=\"user\")],\n", - " cancellation_token=CancellationToken(),\n", - " )\n", - " print(response.inner_messages)\n", - " print(response.chat_message)\n", - "\n", - "\n", - "# Use asyncio.run(assistant_run()) when running in a script.\n", - "await assistant_run()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The call to the {py:meth}`~autogen_agentchat.agents.AssistantAgent.on_messages` method\n", - "returns a {py:class}`~autogen_agentchat.base.Response`\n", - "that contains the agent's final response in the {py:attr}`~autogen_agentchat.base.Response.chat_message` attribute,\n", - "as well as a list of inner messages in the {py:attr}`~autogen_agentchat.base.Response.inner_messages` attribute,\n", - "which stores the agent's \"thought process\" that led to the final response.\n", - "\n", - "```{note}\n", - "It is important to note that {py:meth}`~autogen_agentchat.agents.AssistantAgent.on_messages`\n", - "will update the internal state of the agent -- it will add the messages to the agent's\n", - "history. So you should call this method with new messages.\n", - "**You should not repeatedly call this method with the same messages or the complete history.**\n", - "```\n", - "\n", - "```{note}\n", - "Unlike in v0.2 AgentChat, the tools are executed by the same agent directly within\n", - "the same call to {py:meth}`~autogen_agentchat.agents.AssistantAgent.on_messages`.\n", - "By default, the agent will return the result of the tool call as the final response.\n", - "```\n", - "\n", - "You can also call the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run` method, which is a convenience method that calls {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages`. \n", - "It follows the same interface as [Teams](./teams.ipynb) and returns a {py:class}`~autogen_agentchat.base.TaskResult` object." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Multi-Modal Input\n", - "\n", - "The {py:class}`~autogen_agentchat.agents.AssistantAgent` can handle multi-modal input\n", - "by providing the input as a {py:class}`~autogen_agentchat.messages.MultiModalMessage`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "" - ], - "text/plain": [ - "" - ] - }, - "execution_count": 9, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "from io import BytesIO\n", - "\n", - "import PIL\n", - "import requests\n", - "from autogen_agentchat.messages import MultiModalMessage\n", - "from autogen_core import Image\n", - "\n", - "# Create a multi-modal message with random image and text.\n", - "pil_image = PIL.Image.open(BytesIO(requests.get(\"https://picsum.photos/300/200\").content))\n", - "img = Image(pil_image)\n", - "multi_modal_message = MultiModalMessage(content=[\"Can you describe the content of this image?\", img], source=\"user\")\n", - "img" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "The image depicts a vintage car, likely from the 1930s or 1940s, with a sleek, classic design. The car seems to be customized or well-maintained, as indicated by its shiny exterior and lowered stance. It has a prominent grille and round headlights. There's a license plate on the front with the text \"FARMER BOY.\" The setting appears to be a street with old-style buildings in the background, suggesting a historical or retro theme.\n" - ] - } - ], - "source": [ - "# Use asyncio.run(...) when running in a script.\n", - "response = await agent.on_messages([multi_modal_message], CancellationToken())\n", - "print(response.chat_message.content)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "You can also use {py:class}`~autogen_agentchat.messages.MultiModalMessage` as a `task`\n", - "input to the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run` method." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Streaming Messages\n", - "\n", - "We can also stream each message as it is generated by the agent by using the\n", - "{py:meth}`~autogen_agentchat.agents.AssistantAgent.on_messages_stream` method,\n", - "and use {py:class}`~autogen_agentchat.ui.Console` to print the messages\n", - "as they appear to the console." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- assistant ----------\n", - "[FunctionCall(id='call_fSp5iTGVm2FKw5NIvfECSqNd', arguments='{\"query\":\"AutoGen information\"}', name='web_search')]\n", - "[Prompt tokens: 61, Completion tokens: 16]\n", - "---------- assistant ----------\n", - "[FunctionExecutionResult(content='AutoGen is a programming framework for building multi-agent applications.', call_id='call_fSp5iTGVm2FKw5NIvfECSqNd')]\n", - "---------- assistant ----------\n", - "AutoGen is a programming framework designed for building multi-agent applications. If you need more detailed information or specific aspects about AutoGen, feel free to ask!\n", - "[Prompt tokens: 93, Completion tokens: 32]\n", - "---------- Summary ----------\n", - "Number of inner messages: 2\n", - "Total prompt tokens: 154\n", - "Total completion tokens: 48\n", - "Duration: 4.30 seconds\n" - ] - } - ], - "source": [ - "async def assistant_run_stream() -> None:\n", - " # Option 1: read each message from the stream (as shown in the previous example).\n", - " # async for message in agent.on_messages_stream(\n", - " # [TextMessage(content=\"Find information on AutoGen\", source=\"user\")],\n", - " # cancellation_token=CancellationToken(),\n", - " # ):\n", - " # print(message)\n", - "\n", - " # Option 2: use Console to print all messages as they appear.\n", - " await Console(\n", - " agent.on_messages_stream(\n", - " [TextMessage(content=\"Find information on AutoGen\", source=\"user\")],\n", - " cancellation_token=CancellationToken(),\n", - " ),\n", - " output_stats=True, # Enable stats printing.\n", - " )\n", - "\n", - "\n", - "# Use asyncio.run(assistant_run_stream()) when running in a script.\n", - "await assistant_run_stream()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The {py:meth}`~autogen_agentchat.agents.AssistantAgent.on_messages_stream` method\n", - "returns an asynchronous generator that yields each inner message generated by the agent,\n", - "with the final item being the response message in the {py:attr}`~autogen_agentchat.base.Response.chat_message` attribute.\n", - "\n", - "From the messages, you can observe that the assistant agent utilized the `web_search` tool to\n", - "gather information and responded based on the search results.\n", - "\n", - "You can also use {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run_stream` to get the same streaming behavior as {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream`. It follows the same interface as [Teams](./teams.ipynb)." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Using Tools\n", - "\n", - "Large Language Models (LLMs) are typically limited to generating text or code responses. \n", - "However, many complex tasks benefit from the ability to use external tools that perform specific actions,\n", - "such as fetching data from APIs or databases.\n", - "\n", - "To address this limitation, modern LLMs can now accept a list of available tool schemas \n", - "(descriptions of tools and their arguments) and generate a tool call message. \n", - "This capability is known as **Tool Calling** or **Function Calling** and \n", - "is becoming a popular pattern in building intelligent agent-based applications.\n", - "Refer to the documentation from [OpenAI](https://platform.openai.com/docs/guides/function-calling) \n", - "and [Anthropic](https://docs.anthropic.com/en/docs/build-with-claude/tool-use) for more information about tool calling in LLMs.\n", - "\n", - "In AgentChat, the {py:class}`~autogen_agentchat.agents.AssistantAgent` can use tools to perform specific actions.\n", - "The `web_search` tool is one such tool that allows the assistant agent to search the web for information.\n", - "A custom tool can be a Python function or a subclass of the {py:class}`~autogen_core.tools.BaseTool`.\n", - "\n", - "```{note}\n", - "For how to use model clients directly with tools, refer to the [Tools](../../core-user-guide/components/tools.ipynb) section\n", - "in the Core User Guide.\n", - "```\n", - "\n", - "By default, when {py:class}`~autogen_agentchat.agents.AssistantAgent` executes a tool,\n", - "it will return the tool's output as a string in {py:class}`~autogen_agentchat.messages.ToolCallSummaryMessage` in its response.\n", - "If your tool does not return a well-formed string in natural language, you\n", - "can add a reflection step to have the model summarize the tool's output,\n", - "by setting the `reflect_on_tool_use=True` parameter in the {py:class}`~autogen_agentchat.agents.AssistantAgent` constructor.\n", - "\n", - "### Built-in Tools\n", - "\n", - "AutoGen Extension provides a set of built-in tools that can be used with the Assistant Agent.\n", - "Head over to the [API documentation](../../../reference/index.md) for all the available tools\n", - "under the `autogen_ext.tools` namespace. For example, you can find the following tools:\n", - "\n", - "- {py:mod}`~autogen_ext.tools.graphrag`: Tools for using GraphRAG index.\n", - "- {py:mod}`~autogen_ext.tools.http`: Tools for making HTTP requests.\n", - "- {py:mod}`~autogen_ext.tools.langchain`: Adaptor for using LangChain tools.\n", - "- {py:mod}`~autogen_ext.tools.mcp`: Tools for using Model Chat Protocol (MCP) servers." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Function Tool\n", - "\n", - "The {py:class}`~autogen_agentchat.agents.AssistantAgent` automatically\n", - "converts a Python function into a {py:class}`~autogen_core.tools.FunctionTool`\n", - "which can be used as a tool by the agent and automatically generates the tool schema\n", - "from the function signature and docstring.\n", - "\n", - "The `web_search_func` tool is an example of a function tool.\n", - "The schema is automatically generated." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{'name': 'web_search_func',\n", - " 'description': 'Find information on the web',\n", - " 'parameters': {'type': 'object',\n", - " 'properties': {'query': {'description': 'query',\n", - " 'title': 'Query',\n", - " 'type': 'string'}},\n", - " 'required': ['query'],\n", - " 'additionalProperties': False},\n", - " 'strict': False}" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "from autogen_core.tools import FunctionTool\n", - "\n", - "\n", - "# Define a tool using a Python function.\n", - "async def web_search_func(query: str) -> str:\n", - " \"\"\"Find information on the web\"\"\"\n", - " return \"AutoGen is a programming framework for building multi-agent applications.\"\n", - "\n", - "\n", - "# This step is automatically performed inside the AssistantAgent if the tool is a Python function.\n", - "web_search_function_tool = FunctionTool(web_search_func, description=\"Find information on the web\")\n", - "# The schema is provided to the model during AssistantAgent's on_messages call.\n", - "web_search_function_tool.schema" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Model Context Protocol Tools\n", - "\n", - "The {py:class}`~autogen_agentchat.agents.AssistantAgent` can also use tools that are\n", - "served from a Model Context Protocol (MCP) server\n", - "using {py:func}`~autogen_ext.tools.mcp.mcp_server_tools`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Seattle, located in Washington state, is the most populous city in the state and a major city in the Pacific Northwest region of the United States. It's known for its vibrant cultural scene, significant economic presence, and rich history. Here are some key points about Seattle from the Wikipedia page:\n", - "\n", - "1. **History and Geography**: Seattle is situated between Puget Sound and Lake Washington, with the Cascade Range to the east and the Olympic Mountains to the west. Its history is deeply rooted in Native American heritage and its development was accelerated with the arrival of settlers in the 19th century. The city was officially incorporated in 1869.\n", - "\n", - "2. **Economy**: Seattle is a major economic hub with a diverse economy anchored by sectors like aerospace, technology, and retail. It's home to influential companies such as Amazon and Starbucks, and has a significant impact on the tech industry due to companies like Microsoft and other technology enterprises in the surrounding area.\n", - "\n", - "3. **Cultural Significance**: Known for its music scene, Seattle was the birthplace of grunge music in the early 1990s. It also boasts significant attractions like the Space Needle, Pike Place Market, and the Seattle Art Museum. \n", - "\n", - "4. **Education and Innovation**: The city hosts important educational institutions, with the University of Washington being a leading research university. Seattle is recognized for fostering innovation and is a leader in environmental sustainability efforts.\n", - "\n", - "5. **Demographics and Diversity**: Seattle is noted for its diverse population, reflected in its rich cultural tapestry. It has seen a significant increase in population, leading to urban development and changes in its social landscape.\n", - "\n", - "These points highlight Seattle as a dynamic city with a significant cultural, economic, and educational influence within the United States and beyond.\n" - ] - } - ], - "source": [ - "from autogen_agentchat.agents import AssistantAgent\n", - "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", - "from autogen_ext.tools.mcp import StdioServerParams, mcp_server_tools\n", - "\n", - "# Get the fetch tool from mcp-server-fetch.\n", - "fetch_mcp_server = StdioServerParams(command=\"uvx\", args=[\"mcp-server-fetch\"])\n", - "tools = await mcp_server_tools(fetch_mcp_server)\n", - "\n", - "# Create an agent that can use the fetch tool.\n", - "model_client = OpenAIChatCompletionClient(model=\"gpt-4o\")\n", - "agent = AssistantAgent(name=\"fetcher\", model_client=model_client, tools=tools, reflect_on_tool_use=True) # type: ignore\n", - "\n", - "# Let the agent fetch the content of a URL and summarize it.\n", - "result = await agent.run(task=\"Summarize the content of https://en.wikipedia.org/wiki/Seattle\")\n", - "print(result.messages[-1].content)\n", - "\n", - "# Close the connection to the model client.\n", - "await model_client.close()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Langchain Tools\n", - "\n", - "You can also use tools from the Langchain library\n", - "by wrapping them in {py:class}`~autogen_ext.tools.langchain.LangChainToolAdapter`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- assistant ----------\n", - "[FunctionCall(id='call_BEYRkf53nBS1G2uG60wHP0zf', arguments='{\"query\":\"df[\\'Age\\'].mean()\"}', name='python_repl_ast')]\n", - "[Prompt tokens: 111, Completion tokens: 22]\n", - "---------- assistant ----------\n", - "[FunctionExecutionResult(content='29.69911764705882', call_id='call_BEYRkf53nBS1G2uG60wHP0zf')]\n", - "---------- assistant ----------\n", - "29.69911764705882\n", - "---------- Summary ----------\n", - "Number of inner messages: 2\n", - "Total prompt tokens: 111\n", - "Total completion tokens: 22\n", - "Duration: 0.62 seconds\n" - ] - }, - { - "data": { - "text/plain": [ - "Response(chat_message=ToolCallSummaryMessage(source='assistant', models_usage=None, content='29.69911764705882', type='ToolCallSummaryMessage'), inner_messages=[ToolCallRequestEvent(source='assistant', models_usage=RequestUsage(prompt_tokens=111, completion_tokens=22), content=[FunctionCall(id='call_BEYRkf53nBS1G2uG60wHP0zf', arguments='{\"query\":\"df[\\'Age\\'].mean()\"}', name='python_repl_ast')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='assistant', models_usage=None, content=[FunctionExecutionResult(content='29.69911764705882', call_id='call_BEYRkf53nBS1G2uG60wHP0zf')], type='ToolCallExecutionEvent')])" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "import pandas as pd\n", - "from autogen_ext.tools.langchain import LangChainToolAdapter\n", - "from langchain_experimental.tools.python.tool import PythonAstREPLTool\n", - "\n", - "df = pd.read_csv(\"https://raw.githubusercontent.com/pandas-dev/pandas/main/doc/data/titanic.csv\")\n", - "tool = LangChainToolAdapter(PythonAstREPLTool(locals={\"df\": df}))\n", - "model_client = OpenAIChatCompletionClient(model=\"gpt-4o\")\n", - "agent = AssistantAgent(\n", - " \"assistant\", tools=[tool], model_client=model_client, system_message=\"Use the `df` variable to access the dataset.\"\n", - ")\n", - "await Console(\n", - " agent.on_messages_stream(\n", - " [TextMessage(content=\"What's the average age of the passengers?\", source=\"user\")], CancellationToken()\n", - " ),\n", - " output_stats=True,\n", - ")\n", - "\n", - "await model_client.close()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Parallel Tool Calls\n", - "\n", - "Some models support parallel tool calls, which can be useful for tasks that require multiple tools to be called simultaneously.\n", - "By default, if the model client produces multiple tool calls, {py:class}`~autogen_agentchat.agents.AssistantAgent`\n", - "will call the tools in parallel.\n", - "\n", - "You may want to disable parallel tool calls when the tools have side effects that may interfere with each other, or,\n", - "when agent behavior needs to be consistent across different models.\n", - "This should be done at the model client level.\n", - "\n", - "For {py:class}`~autogen_ext.models.openai.OpenAIChatCompletionClient` and {py:class}`~autogen_ext.models.openai.AzureOpenAIChatCompletionClient`,\n", - "set `parallel_tool_calls=False` to disable parallel tool calls." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "model_client_no_parallel_tool_call = OpenAIChatCompletionClient(\n", - " model=\"gpt-4o\",\n", - " parallel_tool_calls=False, # type: ignore\n", - ")\n", - "agent_no_parallel_tool_call = AssistantAgent(\n", - " name=\"assistant\",\n", - " model_client=model_client_no_parallel_tool_call,\n", - " tools=[web_search],\n", - " system_message=\"Use tools to solve tasks.\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Running an Agent in a Loop\n", - "\n", - "The {py:class}`~autogen_agentchat.agents.AssistantAgent` executes one\n", - "step at a time: one model call, followed by one tool call (or parallel tool calls), and then\n", - "an optional reflection.\n", - "\n", - "To run it in a loop, for example, running it until it stops producing\n", - "tool calls, please refer to [Single-Agent Team](./teams.ipynb#single-agent-team)." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Structured Output\n", - "\n", - "Structured output allows models to return structured JSON text with pre-defined schema\n", - "provided by the application. Different from JSON-mode, the schema can be provided\n", - "as a [Pydantic BaseModel](https://docs.pydantic.dev/latest/concepts/models/)\n", - "class, which can also be used to validate the output. \n", - "\n", - "```{note}\n", - "Structured output is only available for models that support it. It also\n", - "requires the model client to support structured output as well.\n", - "Currently, the {py:class}`~autogen_ext.models.openai.OpenAIChatCompletionClient`\n", - "and {py:class}`~autogen_ext.models.openai.AzureOpenAIChatCompletionClient`\n", - "support structured output.\n", - "```\n", - "\n", - "Structured output is also useful for incorporating Chain-of-Thought\n", - "reasoning in the agent's responses.\n", - "See the example below for how to use structured output with the assistant agent." - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "I am happy.\n", - "---------- assistant ----------\n", - "{\"thoughts\":\"The user explicitly states that they are happy.\",\"response\":\"happy\"}\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='I am happy.', type='TextMessage'), TextMessage(source='assistant', models_usage=RequestUsage(prompt_tokens=89, completion_tokens=18), content='{\"thoughts\":\"The user explicitly states that they are happy.\",\"response\":\"happy\"}', type='TextMessage')], stop_reason=None)" - ] - }, - "execution_count": 2, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "from typing import Literal\n", - "\n", - "from pydantic import BaseModel\n", - "\n", - "\n", - "# The response format for the agent as a Pydantic base model.\n", - "class AgentResponse(BaseModel):\n", - " thoughts: str\n", - " response: Literal[\"happy\", \"sad\", \"neutral\"]\n", - "\n", - "\n", - "# Create an agent that uses the OpenAI GPT-4o model with the custom response format.\n", - "model_client = OpenAIChatCompletionClient(\n", - " model=\"gpt-4o\",\n", - " response_format=AgentResponse, # type: ignore\n", - ")\n", - "agent = AssistantAgent(\n", - " \"assistant\",\n", - " model_client=model_client,\n", - " system_message=\"Categorize the input as happy, sad, or neutral following the JSON format.\",\n", - ")\n", - "\n", - "await Console(agent.run_stream(task=\"I am happy.\"))\n", - "await model_client.close()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Streaming Tokens\n", - "\n", - "You can stream the tokens generated by the model client by setting `model_client_stream=True`.\n", - "This will cause the agent to yield {py:class}`~autogen_agentchat.messages.ModelClientStreamingChunkEvent` messages\n", - "in {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream` and {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run_stream`.\n", - "\n", - "The underlying model API must support streaming tokens for this to work.\n", - "Please check with your model provider to see if this is supported." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "source='assistant' models_usage=None content='Two' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' cities' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' in' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' South' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' America' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' are' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' Buenos' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' Aires' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' in' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' Argentina' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' and' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' São' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' Paulo' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' in' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' Brazil' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content='.' type='ModelClientStreamingChunkEvent'\n", - "Response(chat_message=TextMessage(source='assistant', models_usage=RequestUsage(prompt_tokens=0, completion_tokens=0), content='Two cities in South America are Buenos Aires in Argentina and São Paulo in Brazil.', type='TextMessage'), inner_messages=[])\n" - ] - } - ], - "source": [ - "model_client = OpenAIChatCompletionClient(model=\"gpt-4o\")\n", - "\n", - "streaming_assistant = AssistantAgent(\n", - " name=\"assistant\",\n", - " model_client=model_client,\n", - " system_message=\"You are a helpful assistant.\",\n", - " model_client_stream=True, # Enable streaming tokens.\n", - ")\n", - "\n", - "# Use an async function and asyncio.run() in a script.\n", - "async for message in streaming_assistant.on_messages_stream( # type: ignore\n", - " [TextMessage(content=\"Name two cities in South America\", source=\"user\")],\n", - " cancellation_token=CancellationToken(),\n", - "):\n", - " print(message)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "You can see the streaming chunks in the output above.\n", - "The chunks are generated by the model client and are yielded by the agent as they are received.\n", - "The final response, the concatenation of all the chunks, is yielded right after the last chunk.\n", - "\n", - "Similarly, {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run_stream` will also yield the same streaming chunks,\n", - "followed by a full text message right after the last chunk." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "source='user' models_usage=None content='Name two cities in North America.' type='TextMessage'\n", - "source='assistant' models_usage=None content='Two' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' cities' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' in' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' North' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' America' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' are' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' New' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' York' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' City' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' in' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' the' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' United' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' States' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' and' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' Toronto' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' in' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content=' Canada' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=None content='.' type='ModelClientStreamingChunkEvent'\n", - "source='assistant' models_usage=RequestUsage(prompt_tokens=0, completion_tokens=0) content='Two cities in North America are New York City in the United States and Toronto in Canada.' type='TextMessage'\n", - "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Name two cities in North America.', type='TextMessage'), TextMessage(source='assistant', models_usage=RequestUsage(prompt_tokens=0, completion_tokens=0), content='Two cities in North America are New York City in the United States and Toronto in Canada.', type='TextMessage')], stop_reason=None)\n" - ] - } - ], - "source": [ - "async for message in streaming_assistant.run_stream(task=\"Name two cities in North America.\"): # type: ignore\n", - " print(message)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Using Model Context\n", - "\n", - "{py:class}`~autogen_agentchat.agents.AssistantAgent` has a `model_context`\n", - "parameter that can be used to pass in a {py:class}`~autogen_core.model_context.ChatCompletionContext`\n", - "object. This allows the agent to use different model contexts, such as\n", - "{py:class}`~autogen_core.model_context.BufferedChatCompletionContext` to\n", - "limit the context sent to the model.\n", - "\n", - "By default, {py:class}`~autogen_agentchat.agents.AssistantAgent` uses\n", - "the {py:class}`~autogen_core.model_context.UnboundedChatCompletionContext`\n", - "which sends the full conversation history to the model. To limit the context\n", - "to the last `n` messages, you can use the {py:class}`~autogen_core.model_context.BufferedChatCompletionContext`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from autogen_core.model_context import BufferedChatCompletionContext\n", - "\n", - "# Create an agent that uses only the last 5 messages in the context to generate responses.\n", - "agent = AssistantAgent(\n", - " name=\"assistant\",\n", - " model_client=model_client,\n", - " tools=[web_search],\n", - " system_message=\"Use tools to solve tasks.\",\n", - " model_context=BufferedChatCompletionContext(buffer_size=5), # Only use the last 5 messages in the context.\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Other Preset Agents\n", - "\n", - "The following preset agents are available:\n", - "\n", - "- {py:class}`~autogen_agentchat.agents.UserProxyAgent`: An agent that takes user input returns it as responses.\n", - "- {py:class}`~autogen_agentchat.agents.CodeExecutorAgent`: An agent that can execute code.\n", - "- {py:class}`~autogen_ext.agents.openai.OpenAIAssistantAgent`: An agent that is backed by an OpenAI Assistant, with ability to use custom tools.\n", - "- {py:class}`~autogen_ext.agents.web_surfer.MultimodalWebSurfer`: A multi-modal agent that can search the web and visit web pages for information.\n", - "- {py:class}`~autogen_ext.agents.file_surfer.FileSurfer`: An agent that can search and browse local files for information.\n", - "- {py:class}`~autogen_ext.agents.video_surfer.VideoSurfer`: An agent that can watch videos for information." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Next Step\n", - "\n", - "Having explored the usage of the {py:class}`~autogen_agentchat.agents.AssistantAgent`, we can now proceed to the next section to learn about the teams feature in AgentChat.\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.3" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/messages.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/messages.ipynb deleted file mode 100644 index a3f5a25519dd..000000000000 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/messages.ipynb +++ /dev/null @@ -1,129 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Messages" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In AutoGen AgentChat, _messages_ facilitate communication and information exchange with other agents, orchestrators, and applications. AgentChat supports various message types, each designed for specific purposes." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Types of Messages\n", - "\n", - "At a high level, messages in AgentChat can be categorized into two types: agent-agent messages and an agent's internal events and messages.\n", - "\n", - "### Agent-Agent Messages\n", - "AgentChat supports many message types for agent-to-agent communication. They belong to the union type {py:class}`~autogen_agentchat.messages.ChatMessage`. This message type allows both text and multimodal communication and subsumes other message types, such as {py:class}`~autogen_agentchat.messages.TextMessage` or {py:class}`~autogen_agentchat.messages.MultiModalMessage`.\n", - "\n", - "For example, the following code snippet demonstrates how to create a text message, which accepts a string content and a string source:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from autogen_agentchat.messages import TextMessage\n", - "\n", - "text_message = TextMessage(content=\"Hello, world!\", source=\"User\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Similarly, the following code snippet demonstrates how to create a multimodal message, which accepts\n", - "a list of strings or {py:class}`~autogen_core.Image` objects:" - ] - }, - { - "cell_type": "code", - "execution_count": 16, - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "" - ], - "text/plain": [ - "" - ] - }, - "execution_count": 16, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "from io import BytesIO\n", - "\n", - "import requests\n", - "from autogen_agentchat.messages import MultiModalMessage\n", - "from autogen_core import Image as AGImage\n", - "from PIL import Image\n", - "\n", - "pil_image = Image.open(BytesIO(requests.get(\"https://picsum.photos/300/200\").content))\n", - "img = AGImage(pil_image)\n", - "multi_modal_message = MultiModalMessage(content=[\"Can you describe the content of this image?\", img], source=\"User\")\n", - "img" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The {py:class}`~autogen_agentchat.messages.TextMessage` and {py:class}`~autogen_agentchat.messages.MultiModalMessage` we have created can be passed to agents directly via the {py:class}`~autogen_agentchat.base.ChatAgent.on_messages` method, or as tasks given to a team {py:meth}`~autogen_agentchat.teams.BaseGroupChat.run` method. Messages are also used in the responses of an agent. We will explain these in more detail in [Agents](./agents.ipynb) and [Teams](./teams.ipynb)." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Internal Events\n", - "\n", - "AgentChat also supports the concept of `events` - messages that are internal to an agent. These messages are used to communicate events and information on actions _within_ the agent itself, and belong to the union type {py:class}`~autogen_agentchat.messages.AgentEvent`.\n", - "\n", - "Examples of these include {py:class}`~autogen_agentchat.messages.ToolCallRequestEvent`, which indicates that a request was made to call a tool, and {py:class}`~autogen_agentchat.messages.ToolCallExecutionEvent`, which contains the results of tool calls.\n", - "\n", - "Typically, events are created by the agent itself and are contained in the {py:attr}`~autogen_agentchat.base.Response.inner_messages` field of the {py:class}`~autogen_agentchat.base.Response` returned from {py:class}`~autogen_agentchat.base.ChatAgent.on_messages`. If you are building a custom agent and have events that you want to communicate to other entities (e.g., a UI), you can include these in the {py:attr}`~autogen_agentchat.base.Response.inner_messages` field of the {py:class}`~autogen_agentchat.base.Response`. We will show examples of this in [Custom Agents](../custom-agents.ipynb).\n", - "\n", - "\n", - "You can read about the full set of messages supported in AgentChat in the {py:mod}`~autogen_agentchat.messages` module. \n", - "" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "agnext", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.9" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/state.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/state.ipynb deleted file mode 100644 index 5fd628ac1dd2..000000000000 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/state.ipynb +++ /dev/null @@ -1,359 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Managing State \n", - "\n", - "So far, we have discussed how to build components in a multi-agent application - agents, teams, termination conditions. In many cases, it is useful to save the state of these components to disk and load them back later. This is particularly useful in a web application where stateless endpoints respond to requests and need to load the state of the application from persistent storage.\n", - "\n", - "In this notebook, we will discuss how to save and load the state of agents, teams, and termination conditions. \n", - " \n", - "\n", - "## Saving and Loading Agents\n", - "\n", - "We can get the state of an agent by calling {py:meth}`~autogen_agentchat.agents.AssistantAgent.save_state` method on \n", - "an {py:class}`~autogen_agentchat.agents.AssistantAgent`. " - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "In Tanganyika's embrace so wide and deep, \n", - "Ancient waters cradle secrets they keep, \n", - "Echoes of time where horizons sleep. \n" - ] - } - ], - "source": [ - "from autogen_agentchat.agents import AssistantAgent\n", - "from autogen_agentchat.conditions import MaxMessageTermination\n", - "from autogen_agentchat.messages import TextMessage\n", - "from autogen_agentchat.teams import RoundRobinGroupChat\n", - "from autogen_agentchat.ui import Console\n", - "from autogen_core import CancellationToken\n", - "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", - "\n", - "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-2024-08-06\")\n", - "\n", - "assistant_agent = AssistantAgent(\n", - " name=\"assistant_agent\",\n", - " system_message=\"You are a helpful assistant\",\n", - " model_client=model_client,\n", - ")\n", - "\n", - "# Use asyncio.run(...) when running in a script.\n", - "response = await assistant_agent.on_messages(\n", - " [TextMessage(content=\"Write a 3 line poem on lake tangayika\", source=\"user\")], CancellationToken()\n", - ")\n", - "print(response.chat_message.content)\n", - "await model_client.close()" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "{'type': 'AssistantAgentState', 'version': '1.0.0', 'llm_messages': [{'content': 'Write a 3 line poem on lake tangayika', 'source': 'user', 'type': 'UserMessage'}, {'content': \"In Tanganyika's embrace so wide and deep, \\nAncient waters cradle secrets they keep, \\nEchoes of time where horizons sleep. \", 'source': 'assistant_agent', 'type': 'AssistantMessage'}]}\n" - ] - } - ], - "source": [ - "agent_state = await assistant_agent.save_state()\n", - "print(agent_state)" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "The last line of the poem was: \"Echoes of time where horizons sleep.\"\n" - ] - } - ], - "source": [ - "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-2024-08-06\")\n", - "\n", - "new_assistant_agent = AssistantAgent(\n", - " name=\"assistant_agent\",\n", - " system_message=\"You are a helpful assistant\",\n", - " model_client=model_client,\n", - ")\n", - "await new_assistant_agent.load_state(agent_state)\n", - "\n", - "# Use asyncio.run(...) when running in a script.\n", - "response = await new_assistant_agent.on_messages(\n", - " [TextMessage(content=\"What was the last line of the previous poem you wrote\", source=\"user\")], CancellationToken()\n", - ")\n", - "print(response.chat_message.content)\n", - "await model_client.close()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{note}\n", - "For {py:class}`~autogen_agentchat.agents.AssistantAgent`, its state consists of the model_context.\n", - "If your write your own custom agent, consider overriding the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.save_state` and {py:meth}`~autogen_agentchat.agents.BaseChatAgent.load_state` methods to customize the behavior. The default implementations save and load an empty state.\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Saving and Loading Teams \n", - "\n", - "We can get the state of a team by calling `save_state` method on the team and load it back by calling `load_state` method on the team. \n", - "\n", - "When we call `save_state` on a team, it saves the state of all the agents in the team.\n", - "\n", - "We will begin by creating a simple {py:class}`~autogen_agentchat.teams.RoundRobinGroupChat` team with a single agent and ask it to write a poem. " - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "Write a beautiful poem 3-line about lake tangayika\n", - "---------- assistant_agent ----------\n", - "In Tanganyika's gleam, beneath the azure skies, \n", - "Whispers of ancient waters, in tranquil guise, \n", - "Nature's mirror, where dreams and serenity lie.\n", - "[Prompt tokens: 29, Completion tokens: 34]\n", - "---------- Summary ----------\n", - "Number of messages: 2\n", - "Finish reason: Maximum number of messages 2 reached, current message count: 2\n", - "Total prompt tokens: 29\n", - "Total completion tokens: 34\n", - "Duration: 0.71 seconds\n" - ] - } - ], - "source": [ - "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-2024-08-06\")\n", - "\n", - "# Define a team.\n", - "assistant_agent = AssistantAgent(\n", - " name=\"assistant_agent\",\n", - " system_message=\"You are a helpful assistant\",\n", - " model_client=model_client,\n", - ")\n", - "agent_team = RoundRobinGroupChat([assistant_agent], termination_condition=MaxMessageTermination(max_messages=2))\n", - "\n", - "# Run the team and stream messages to the console.\n", - "stream = agent_team.run_stream(task=\"Write a beautiful poem 3-line about lake tangayika\")\n", - "\n", - "# Use asyncio.run(...) when running in a script.\n", - "await Console(stream)\n", - "\n", - "# Save the state of the agent team.\n", - "team_state = await agent_team.save_state()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If we reset the team (simulating instantiation of the team), and ask the question `What was the last line of the poem you wrote?`, we see that the team is unable to accomplish this as there is no reference to the previous run." - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "What was the last line of the poem you wrote?\n", - "---------- assistant_agent ----------\n", - "I'm sorry, but I am unable to recall or access previous interactions, including any specific poem I may have composed in our past conversations. If you like, I can write a new poem for you.\n", - "[Prompt tokens: 28, Completion tokens: 40]\n", - "---------- Summary ----------\n", - "Number of messages: 2\n", - "Finish reason: Maximum number of messages 2 reached, current message count: 2\n", - "Total prompt tokens: 28\n", - "Total completion tokens: 40\n", - "Duration: 0.70 seconds\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What was the last line of the poem you wrote?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=28, completion_tokens=40), content=\"I'm sorry, but I am unable to recall or access previous interactions, including any specific poem I may have composed in our past conversations. If you like, I can write a new poem for you.\", type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')" - ] - }, - "execution_count": 6, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "await agent_team.reset()\n", - "stream = agent_team.run_stream(task=\"What was the last line of the poem you wrote?\")\n", - "await Console(stream)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Next, we load the state of the team and ask the same question. We see that the team is able to accurately return the last line of the poem it wrote." - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "{'type': 'TeamState', 'version': '1.0.0', 'agent_states': {'group_chat_manager/a55364ad-86fd-46ab-9449-dcb5260b1e06': {'type': 'RoundRobinManagerState', 'version': '1.0.0', 'message_thread': [{'source': 'user', 'models_usage': None, 'content': 'Write a beautiful poem 3-line about lake tangayika', 'type': 'TextMessage'}, {'source': 'assistant_agent', 'models_usage': {'prompt_tokens': 29, 'completion_tokens': 34}, 'content': \"In Tanganyika's gleam, beneath the azure skies, \\nWhispers of ancient waters, in tranquil guise, \\nNature's mirror, where dreams and serenity lie.\", 'type': 'TextMessage'}], 'current_turn': 0, 'next_speaker_index': 0}, 'collect_output_messages/a55364ad-86fd-46ab-9449-dcb5260b1e06': {}, 'assistant_agent/a55364ad-86fd-46ab-9449-dcb5260b1e06': {'type': 'ChatAgentContainerState', 'version': '1.0.0', 'agent_state': {'type': 'AssistantAgentState', 'version': '1.0.0', 'llm_messages': [{'content': 'Write a beautiful poem 3-line about lake tangayika', 'source': 'user', 'type': 'UserMessage'}, {'content': \"In Tanganyika's gleam, beneath the azure skies, \\nWhispers of ancient waters, in tranquil guise, \\nNature's mirror, where dreams and serenity lie.\", 'source': 'assistant_agent', 'type': 'AssistantMessage'}]}, 'message_buffer': []}}, 'team_id': 'a55364ad-86fd-46ab-9449-dcb5260b1e06'}\n", - "---------- user ----------\n", - "What was the last line of the poem you wrote?\n", - "---------- assistant_agent ----------\n", - "The last line of the poem I wrote is: \n", - "\"Nature's mirror, where dreams and serenity lie.\"\n", - "[Prompt tokens: 86, Completion tokens: 22]\n", - "---------- Summary ----------\n", - "Number of messages: 2\n", - "Finish reason: Maximum number of messages 2 reached, current message count: 2\n", - "Total prompt tokens: 86\n", - "Total completion tokens: 22\n", - "Duration: 0.96 seconds\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What was the last line of the poem you wrote?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=22), content='The last line of the poem I wrote is: \\n\"Nature\\'s mirror, where dreams and serenity lie.\"', type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')" - ] - }, - "execution_count": 7, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "print(team_state)\n", - "\n", - "# Load team state.\n", - "await agent_team.load_state(team_state)\n", - "stream = agent_team.run_stream(task=\"What was the last line of the poem you wrote?\")\n", - "await Console(stream)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Persisting State (File or Database)\n", - "\n", - "In many cases, we may want to persist the state of the team to disk (or a database) and load it back later. State is a dictionary that can be serialized to a file or written to a database." - ] - }, - { - "cell_type": "code", - "execution_count": 13, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "What was the last line of the poem you wrote?\n", - "---------- assistant_agent ----------\n", - "The last line of the poem I wrote is: \n", - "\"Nature's mirror, where dreams and serenity lie.\"\n", - "[Prompt tokens: 86, Completion tokens: 22]\n", - "---------- Summary ----------\n", - "Number of messages: 2\n", - "Finish reason: Maximum number of messages 2 reached, current message count: 2\n", - "Total prompt tokens: 86\n", - "Total completion tokens: 22\n", - "Duration: 0.72 seconds\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What was the last line of the poem you wrote?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=22), content='The last line of the poem I wrote is: \\n\"Nature\\'s mirror, where dreams and serenity lie.\"', type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')" - ] - }, - "execution_count": 13, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "import json\n", - "\n", - "## save state to disk\n", - "\n", - "with open(\"coding/team_state.json\", \"w\") as f:\n", - " json.dump(team_state, f)\n", - "\n", - "## load state from disk\n", - "with open(\"coding/team_state.json\", \"r\") as f:\n", - " team_state = json.load(f)\n", - "\n", - "new_agent_team = RoundRobinGroupChat([assistant_agent], termination_condition=MaxMessageTermination(max_messages=2))\n", - "await new_agent_team.load_state(team_state)\n", - "stream = new_agent_team.run_stream(task=\"What was the last line of the poem you wrote?\")\n", - "await Console(stream)\n", - "await model_client.close()" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "agnext", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.9" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/termination.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/termination.ipynb deleted file mode 100644 index 8c9701550dad..000000000000 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/termination.ipynb +++ /dev/null @@ -1,517 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Termination \n", - "\n", - "In the previous section, we explored how to define agents, and organize them into teams that can solve tasks. However, a run can go on forever, and in many cases, we need to know _when_ to stop them. This is the role of the termination condition.\n", - "\n", - "AgentChat supports several termination condition by providing a base {py:class}`~autogen_agentchat.base.TerminationCondition` class and several implementations that inherit from it.\n", - "\n", - "A termination condition is a callable that takes a sequence of {py:class}`~autogen_agentchat.messages.AgentEvent` or {py:class}`~autogen_agentchat.messages.ChatMessage` objects **since the last time the condition was called**, and returns a {py:class}`~autogen_agentchat.messages.StopMessage` if the conversation should be terminated, or `None` otherwise.\n", - "Once a termination condition has been reached, it must be reset by calling {py:meth}`~autogen_agentchat.base.TerminationCondition.reset` before it can be used again.\n", - "\n", - "Some important things to note about termination conditions: \n", - "- They are stateful but reset automatically after each run ({py:meth}`~autogen_agentchat.base.TaskRunner.run` or {py:meth}`~autogen_agentchat.base.TaskRunner.run_stream`) is finished.\n", - "- They can be combined using the AND and OR operators.\n", - "\n", - "```{note}\n", - "For group chat teams (i.e., {py:class}`~autogen_agentchat.teams.RoundRobinGroupChat`,\n", - "{py:class}`~autogen_agentchat.teams.SelectorGroupChat`, and {py:class}`~autogen_agentchat.teams.Swarm`),\n", - "the termination condition is called after each agent responds.\n", - "While a response may contain multiple inner messages, the team calls its termination condition just once for all the messages from a single response.\n", - "So the condition is called with the \"delta sequence\" of messages since the last time it was called.\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Built-In Termination Conditions: \n", - "1. {py:class}`~autogen_agentchat.conditions.MaxMessageTermination`: Stops after a specified number of messages have been produced, including both agent and task messages.\n", - "2. {py:class}`~autogen_agentchat.conditions.TextMentionTermination`: Stops when specific text or string is mentioned in a message (e.g., \"TERMINATE\").\n", - "3. {py:class}`~autogen_agentchat.conditions.TokenUsageTermination`: Stops when a certain number of prompt or completion tokens are used. This requires the agents to report token usage in their messages.\n", - "4. {py:class}`~autogen_agentchat.conditions.TimeoutTermination`: Stops after a specified duration in seconds.\n", - "5. {py:class}`~autogen_agentchat.conditions.HandoffTermination`: Stops when a handoff to a specific target is requested. Handoff messages can be used to build patterns such as {py:class}`~autogen_agentchat.teams.Swarm`. This is useful when you want to pause the run and allow application or user to provide input when an agent hands off to them.\n", - "6. {py:class}`~autogen_agentchat.conditions.SourceMatchTermination`: Stops after a specific agent responds.\n", - "7. {py:class}`~autogen_agentchat.conditions.ExternalTermination`: Enables programmatic control of termination from outside the run. This is useful for UI integration (e.g., \"Stop\" buttons in chat interfaces).\n", - "8. {py:class}`~autogen_agentchat.conditions.StopMessageTermination`: Stops when a {py:class}`~autogen_agentchat.messages.StopMessage` is produced by an agent.\n", - "9. {py:class}`~autogen_agentchat.conditions.TextMessageTermination`: Stops when a {py:class}`~autogen_agentchat.messages.TextMessage` is produced by an agent.\n", - "10. {py:class}`~autogen_agentchat.conditions.FunctionCallTermination`: Stops when a {py:class}`~autogen_agentchat.messages.ToolCallExecutionEvent` containing a {py:class}`~autogen_core.models.FunctionExecutionResult` with a matching name is produced by an agent." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Basic Usage\n", - "\n", - "To demonstrate the characteristics of termination conditions, we'll create a team consisting of two agents: a primary agent responsible for text generation and a critic agent that reviews and provides feedback on the generated text." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from autogen_agentchat.agents import AssistantAgent\n", - "from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination\n", - "from autogen_agentchat.teams import RoundRobinGroupChat\n", - "from autogen_agentchat.ui import Console\n", - "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", - "\n", - "model_client = OpenAIChatCompletionClient(\n", - " model=\"gpt-4o\",\n", - " temperature=1,\n", - " # api_key=\"sk-...\", # Optional if you have an OPENAI_API_KEY env variable set.\n", - ")\n", - "\n", - "# Create the primary agent.\n", - "primary_agent = AssistantAgent(\n", - " \"primary\",\n", - " model_client=model_client,\n", - " system_message=\"You are a helpful AI assistant.\",\n", - ")\n", - "\n", - "# Create the critic agent.\n", - "critic_agent = AssistantAgent(\n", - " \"critic\",\n", - " model_client=model_client,\n", - " system_message=\"Provide constructive feedback for every message. Respond with 'APPROVE' to when your feedbacks are addressed.\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Let's explore how termination conditions automatically reset after each `run` or `run_stream` call, allowing the team to resume its conversation from where it left off." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "Write a unique, Haiku about the weather in Paris\n", - "---------- primary ----------\n", - "Gentle rain whispers, \n", - "Cobblestones glisten softly— \n", - "Paris dreams in gray.\n", - "[Prompt tokens: 30, Completion tokens: 19]\n", - "---------- critic ----------\n", - "The Haiku captures the essence of a rainy day in Paris beautifully, and the imagery is vivid. However, it's important to ensure the use of the traditional 5-7-5 syllable structure for Haikus. Your current Haiku lines are composed of 4-7-5 syllables, which slightly deviates from the form. Consider revising the first line to fit the structure.\n", - "\n", - "For example:\n", - "Soft rain whispers down, \n", - "Cobblestones glisten softly — \n", - "Paris dreams in gray.\n", - "\n", - "This revision maintains the essence of your original lines while adhering to the traditional Haiku structure.\n", - "[Prompt tokens: 70, Completion tokens: 120]\n", - "---------- Summary ----------\n", - "Number of messages: 3\n", - "Finish reason: Maximum number of messages 3 reached, current message count: 3\n", - "Total prompt tokens: 100\n", - "Total completion tokens: 139\n", - "Duration: 3.34 seconds\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Write a unique, Haiku about the weather in Paris'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=30, completion_tokens=19), content='Gentle rain whispers, \\nCobblestones glisten softly— \\nParis dreams in gray.'), TextMessage(source='critic', models_usage=RequestUsage(prompt_tokens=70, completion_tokens=120), content=\"The Haiku captures the essence of a rainy day in Paris beautifully, and the imagery is vivid. However, it's important to ensure the use of the traditional 5-7-5 syllable structure for Haikus. Your current Haiku lines are composed of 4-7-5 syllables, which slightly deviates from the form. Consider revising the first line to fit the structure.\\n\\nFor example:\\nSoft rain whispers down, \\nCobblestones glisten softly — \\nParis dreams in gray.\\n\\nThis revision maintains the essence of your original lines while adhering to the traditional Haiku structure.\")], stop_reason='Maximum number of messages 3 reached, current message count: 3')" - ] - }, - "execution_count": 4, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "max_msg_termination = MaxMessageTermination(max_messages=3)\n", - "round_robin_team = RoundRobinGroupChat([primary_agent, critic_agent], termination_condition=max_msg_termination)\n", - "\n", - "# Use asyncio.run(...) if you are running this script as a standalone script.\n", - "await Console(round_robin_team.run_stream(task=\"Write a unique, Haiku about the weather in Paris\"))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The conversation stopped after reaching the maximum message limit. Since the primary agent didn't get to respond to the feedback, let's continue the conversation." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- primary ----------\n", - "Thank you for your feedback. Here is the revised Haiku:\n", - "\n", - "Soft rain whispers down, \n", - "Cobblestones glisten softly — \n", - "Paris dreams in gray.\n", - "[Prompt tokens: 181, Completion tokens: 32]\n", - "---------- critic ----------\n", - "The revised Haiku now follows the traditional 5-7-5 syllable pattern, and it still beautifully captures the atmospheric mood of Paris in the rain. The imagery and flow are both clear and evocative. Well done on making the adjustment! \n", - "\n", - "APPROVE\n", - "[Prompt tokens: 234, Completion tokens: 54]\n", - "---------- primary ----------\n", - "Thank you for your kind words and approval. I'm glad the revision meets your expectations and captures the essence of Paris. If you have any more requests or need further assistance, feel free to ask!\n", - "[Prompt tokens: 279, Completion tokens: 39]\n", - "---------- Summary ----------\n", - "Number of messages: 3\n", - "Finish reason: Maximum number of messages 3 reached, current message count: 3\n", - "Total prompt tokens: 694\n", - "Total completion tokens: 125\n", - "Duration: 6.43 seconds\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=181, completion_tokens=32), content='Thank you for your feedback. Here is the revised Haiku:\\n\\nSoft rain whispers down, \\nCobblestones glisten softly — \\nParis dreams in gray.'), TextMessage(source='critic', models_usage=RequestUsage(prompt_tokens=234, completion_tokens=54), content='The revised Haiku now follows the traditional 5-7-5 syllable pattern, and it still beautifully captures the atmospheric mood of Paris in the rain. The imagery and flow are both clear and evocative. Well done on making the adjustment! \\n\\nAPPROVE'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=279, completion_tokens=39), content=\"Thank you for your kind words and approval. I'm glad the revision meets your expectations and captures the essence of Paris. If you have any more requests or need further assistance, feel free to ask!\")], stop_reason='Maximum number of messages 3 reached, current message count: 3')" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "# Use asyncio.run(...) if you are running this script as a standalone script.\n", - "await Console(round_robin_team.run_stream())" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The team continued from where it left off, allowing the primary agent to respond to the feedback." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Combining Termination Conditions\n", - "\n", - "Let's show how termination conditions can be combined using the AND (`&`) and OR (`|`) operators to create more complex termination logic. For example, we'll create a team that stops either after 10 messages are generated or when the critic agent approves a message.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "Write a unique, Haiku about the weather in Paris\n", - "---------- primary ----------\n", - "Spring breeze gently hums, \n", - "Cherry blossoms in full bloom— \n", - "Paris wakes to life.\n", - "[Prompt tokens: 467, Completion tokens: 19]\n", - "---------- critic ----------\n", - "The Haiku beautifully captures the awakening of Paris in the spring. The imagery of a gentle spring breeze and cherry blossoms in full bloom effectively conveys the rejuvenating feel of the season. The final line, \"Paris wakes to life,\" encapsulates the renewed energy and vibrancy of the city. The Haiku adheres to the 5-7-5 syllable structure and portrays a vivid seasonal transformation in a concise and poetic manner. Excellent work!\n", - "\n", - "APPROVE\n", - "[Prompt tokens: 746, Completion tokens: 93]\n", - "---------- Summary ----------\n", - "Number of messages: 3\n", - "Finish reason: Text 'APPROVE' mentioned\n", - "Total prompt tokens: 1213\n", - "Total completion tokens: 112\n", - "Duration: 2.75 seconds\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Write a unique, Haiku about the weather in Paris'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=467, completion_tokens=19), content='Spring breeze gently hums, \\nCherry blossoms in full bloom— \\nParis wakes to life.'), TextMessage(source='critic', models_usage=RequestUsage(prompt_tokens=746, completion_tokens=93), content='The Haiku beautifully captures the awakening of Paris in the spring. The imagery of a gentle spring breeze and cherry blossoms in full bloom effectively conveys the rejuvenating feel of the season. The final line, \"Paris wakes to life,\" encapsulates the renewed energy and vibrancy of the city. The Haiku adheres to the 5-7-5 syllable structure and portrays a vivid seasonal transformation in a concise and poetic manner. Excellent work!\\n\\nAPPROVE')], stop_reason=\"Text 'APPROVE' mentioned\")" - ] - }, - "execution_count": 9, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "max_msg_termination = MaxMessageTermination(max_messages=10)\n", - "text_termination = TextMentionTermination(\"APPROVE\")\n", - "combined_termination = max_msg_termination | text_termination\n", - "\n", - "round_robin_team = RoundRobinGroupChat([primary_agent, critic_agent], termination_condition=combined_termination)\n", - "\n", - "# Use asyncio.run(...) if you are running this script as a standalone script.\n", - "await Console(round_robin_team.run_stream(task=\"Write a unique, Haiku about the weather in Paris\"))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The conversation stopped after the critic agent approved the message, although it could have also stopped if 10 messages were generated.\n", - "\n", - "Alternatively, if we want to stop the run only when both conditions are met, we can use the AND (`&`) operator." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "combined_termination = max_msg_termination & text_termination" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Custom Termination Condition\n", - "\n", - "The built-in termination conditions are sufficient for most use cases.\n", - "However, there may be cases where you need to implement a custom termination condition that doesn't fit into the existing ones.\n", - "You can do this by subclassing the {py:class}`~autogen_agentchat.base.TerminationCondition` class.\n", - "\n", - "In this example, we create a custom termination condition that stops the conversation when\n", - "a specific function call is made." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from typing import Sequence\n", - "\n", - "from autogen_agentchat.base import TerminatedException, TerminationCondition\n", - "from autogen_agentchat.messages import AgentEvent, ChatMessage, StopMessage, ToolCallExecutionEvent\n", - "from autogen_core import Component\n", - "from pydantic import BaseModel\n", - "from typing_extensions import Self\n", - "\n", - "\n", - "class FunctionCallTerminationConfig(BaseModel):\n", - " \"\"\"Configuration for the termination condition to allow for serialization\n", - " and deserialization of the component.\n", - " \"\"\"\n", - "\n", - " function_name: str\n", - "\n", - "\n", - "class FunctionCallTermination(TerminationCondition, Component[FunctionCallTerminationConfig]):\n", - " \"\"\"Terminate the conversation if a FunctionExecutionResult with a specific name is received.\"\"\"\n", - "\n", - " component_config_schema = FunctionCallTerminationConfig\n", - " \"\"\"The schema for the component configuration.\"\"\"\n", - "\n", - " def __init__(self, function_name: str) -> None:\n", - " self._terminated = False\n", - " self._function_name = function_name\n", - "\n", - " @property\n", - " def terminated(self) -> bool:\n", - " return self._terminated\n", - "\n", - " async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None:\n", - " if self._terminated:\n", - " raise TerminatedException(\"Termination condition has already been reached\")\n", - " for message in messages:\n", - " if isinstance(message, ToolCallExecutionEvent):\n", - " for execution in message.content:\n", - " if execution.name == self._function_name:\n", - " self._terminated = True\n", - " return StopMessage(\n", - " content=f\"Function '{self._function_name}' was executed.\",\n", - " source=\"FunctionCallTermination\",\n", - " )\n", - " return None\n", - "\n", - " async def reset(self) -> None:\n", - " self._terminated = False\n", - "\n", - " def _to_config(self) -> FunctionCallTerminationConfig:\n", - " return FunctionCallTerminationConfig(\n", - " function_name=self._function_name,\n", - " )\n", - "\n", - " @classmethod\n", - " def _from_config(cls, config: FunctionCallTerminationConfig) -> Self:\n", - " return cls(\n", - " function_name=config.function_name,\n", - " )" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Let's use this new termination condition to stop the conversation when the critic agent approves a message\n", - "using the `approve` function call.\n", - "\n", - "First we create a simple function that will be called when the critic agent approves a message." - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": {}, - "outputs": [], - "source": [ - "def approve() -> None:\n", - " \"\"\"Approve the message when all feedbacks have been addressed.\"\"\"\n", - " pass" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Then we create the agents. The critic agent is equipped with the `approve` tool." - ] - }, - { - "cell_type": "code", - "execution_count": 9, - "metadata": {}, - "outputs": [], - "source": [ - "from autogen_agentchat.agents import AssistantAgent\n", - "from autogen_agentchat.teams import RoundRobinGroupChat\n", - "from autogen_agentchat.ui import Console\n", - "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", - "\n", - "model_client = OpenAIChatCompletionClient(\n", - " model=\"gpt-4o\",\n", - " temperature=1,\n", - " # api_key=\"sk-...\", # Optional if you have an OPENAI_API_KEY env variable set.\n", - ")\n", - "\n", - "# Create the primary agent.\n", - "primary_agent = AssistantAgent(\n", - " \"primary\",\n", - " model_client=model_client,\n", - " system_message=\"You are a helpful AI assistant.\",\n", - ")\n", - "\n", - "# Create the critic agent with the approve function as a tool.\n", - "critic_agent = AssistantAgent(\n", - " \"critic\",\n", - " model_client=model_client,\n", - " tools=[approve], # Register the approve function as a tool.\n", - " system_message=\"Provide constructive feedback. Use the approve tool to approve when all feedbacks are addressed.\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now, we create the termination condition and the team.\n", - "We run the team with the poem-writing task." - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "Write a unique, Haiku about the weather in Paris\n", - "---------- primary ----------\n", - "Raindrops gently fall, \n", - "Cobblestones shine in dim light— \n", - "Paris dreams in grey. \n", - "---------- critic ----------\n", - "This Haiku beautifully captures a melancholic yet romantic image of Paris in the rain. The use of sensory imagery like \"Raindrops gently fall\" and \"Cobblestones shine\" effectively paints a vivid picture. It could be interesting to experiment with more distinct seasonal elements of Paris, such as incorporating the Seine River or iconic landmarks in the context of the weather. Overall, it successfully conveys the atmosphere of Paris in subtle, poetic imagery.\n", - "---------- primary ----------\n", - "Thank you for your feedback! I’m glad you enjoyed the imagery. Here’s another Haiku that incorporates iconic Parisian elements:\n", - "\n", - "Eiffel stands in mist, \n", - "Seine's ripple mirrors the sky— \n", - "Spring whispers anew. \n", - "---------- critic ----------\n", - "[FunctionCall(id='call_QEWJZ873EG4UIEpsQHi1HsAu', arguments='{}', name='approve')]\n", - "---------- critic ----------\n", - "[FunctionExecutionResult(content='None', name='approve', call_id='call_QEWJZ873EG4UIEpsQHi1HsAu', is_error=False)]\n", - "---------- critic ----------\n", - "None\n" - ] - }, - { - "data": { - "text/plain": [ - "TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Write a unique, Haiku about the weather in Paris', type='TextMessage'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=30, completion_tokens=23), metadata={}, content='Raindrops gently fall, \\nCobblestones shine in dim light— \\nParis dreams in grey. ', type='TextMessage'), TextMessage(source='critic', models_usage=RequestUsage(prompt_tokens=99, completion_tokens=90), metadata={}, content='This Haiku beautifully captures a melancholic yet romantic image of Paris in the rain. The use of sensory imagery like \"Raindrops gently fall\" and \"Cobblestones shine\" effectively paints a vivid picture. It could be interesting to experiment with more distinct seasonal elements of Paris, such as incorporating the Seine River or iconic landmarks in the context of the weather. Overall, it successfully conveys the atmosphere of Paris in subtle, poetic imagery.', type='TextMessage'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=152, completion_tokens=48), metadata={}, content=\"Thank you for your feedback! I’m glad you enjoyed the imagery. Here’s another Haiku that incorporates iconic Parisian elements:\\n\\nEiffel stands in mist, \\nSeine's ripple mirrors the sky— \\nSpring whispers anew. \", type='TextMessage'), ToolCallRequestEvent(source='critic', models_usage=RequestUsage(prompt_tokens=246, completion_tokens=11), metadata={}, content=[FunctionCall(id='call_QEWJZ873EG4UIEpsQHi1HsAu', arguments='{}', name='approve')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='critic', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='None', name='approve', call_id='call_QEWJZ873EG4UIEpsQHi1HsAu', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='critic', models_usage=None, metadata={}, content='None', type='ToolCallSummaryMessage')], stop_reason=\"Function 'approve' was executed.\")" - ] - }, - "execution_count": 10, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "function_call_termination = FunctionCallTermination(function_name=\"approve\")\n", - "round_robin_team = RoundRobinGroupChat([primary_agent, critic_agent], termination_condition=function_call_termination)\n", - "\n", - "# Use asyncio.run(...) if you are running this script as a standalone script.\n", - "await Console(round_robin_team.run_stream(task=\"Write a unique, Haiku about the weather in Paris\"))\n", - "await model_client.close()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "You can see that the conversation stopped when the critic agent approved the message using the `approve` function call." - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.7" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/python/packages/autogen-core/pyproject.toml b/python/packages/autogen-core/pyproject.toml index 4733368a3871..5d261d0390c8 100644 --- a/python/packages/autogen-core/pyproject.toml +++ b/python/packages/autogen-core/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "autogen-core" -version = "0.4.9" +version = "0.7.2" license = {file = "LICENSE-CODE"} description = "Foundational interfaces and agent runtime implementation for AutoGen" readme = "README.md" @@ -19,7 +19,7 @@ dependencies = [ "typing-extensions>=4.0.0", "pydantic<3.0.0,>=2.10.0", "protobuf~=5.29.3", - "opentelemetry-api>=1.27.0", + "opentelemetry-api>=1.34.1", "jsonref~=1.1.0", ] @@ -42,7 +42,7 @@ dev = [ "llama-index", "markdownify", "nbqa", - "opentelemetry-sdk>=1.27.0", + "opentelemetry-sdk>=1.34.1", "pip", "polars", "python-dotenv", @@ -57,34 +57,13 @@ dev = [ "types-protobuf", "types-requests", "wikipedia", - - # Documentation - "myst-nb==1.1.2", - "pydata-sphinx-theme==0.16.0", - "sphinx-copybutton", - "sphinx-design", - "sphinx", - "sphinxcontrib-apidoc", - "autodoc_pydantic~=2.2", - "pygments", - "sphinxext-rediraffe", - - "autogen_ext==0.4.9", - - # Documentation tooling - "diskcache", - "redis", - "sphinx-autobuild", ] [tool.ruff] extend = "../../pyproject.toml" exclude = ["build", "dist", "src/autogen_core/application/protos", "tests/protos"] -include = ["src/**", "docs/**/*.ipynb", "tests/**"] - -[tool.ruff.lint.per-file-ignores] -"docs/**.ipynb" = ["T20"] +include = ["src/**", "tests/**"] [tool.pyright] extends = "../../pyproject.toml" @@ -97,11 +76,6 @@ minversion = "6.0" testpaths = ["tests"] asyncio_default_fixture_loop_scope = "session" -[tool.nbqa.addopts] -mypy = [ - "--disable-error-code=top-level-await" -] - [tool.poe] include = "../../shared_tasks.toml" @@ -110,24 +84,5 @@ test = "pytest -n auto --cov=src --cov-report=term-missing --cov-report=xml" mypy.default_item_type = "cmd" mypy.sequence = [ "mypy --config-file ../../pyproject.toml --exclude src/autogen_core/application/protos --exclude tests/protos src tests", - "nbqa mypy docs/src --config-file ../../pyproject.toml", ] -# Docs -docs-clean = "rm -rf docs/build" - -# Inline tables are WAY easier to read but for some reason they break pyright. So we have to write it out this way. -# Example of inline table: -# docs-build = [ -# "docs-apidoc-all", -# { cmd = "sphinx-build docs/src docs/build" } -# ] - -docs-build = "sphinx-build docs/src docs/build" - -docs-serve = "sphinx-autobuild --watch src docs/src docs/build --port 8000 --jobs auto" - -docs-check = "sphinx-build --fail-on-warning docs/src docs/build" - -docs-check-examples = "sphinx-build -b code_lint docs/src docs/build" - diff --git a/python/packages/autogen-core/src/autogen_core/__init__.py b/python/packages/autogen-core/src/autogen_core/__init__.py index 0198544ca61e..ffc8e984aee8 100644 --- a/python/packages/autogen-core/src/autogen_core/__init__.py +++ b/python/packages/autogen-core/src/autogen_core/__init__.py @@ -59,6 +59,11 @@ from ._single_threaded_agent_runtime import SingleThreadedAgentRuntime from ._subscription import Subscription from ._subscription_context import SubscriptionInstantiationContext +from ._telemetry import ( + trace_create_agent_span, + trace_invoke_agent_span, + trace_tool_span, +) from ._topic import TopicId from ._type_prefix_subscription import TypePrefixSubscription from ._type_subscription import TypeSubscription @@ -132,4 +137,7 @@ "DropMessage", "InterventionHandler", "DefaultInterventionHandler", + "trace_create_agent_span", + "trace_invoke_agent_span", + "trace_tool_span", ] diff --git a/python/packages/autogen-core/src/autogen_core/_agent.py b/python/packages/autogen-core/src/autogen_core/_agent.py index 0f37b822ff8a..e407fe137394 100644 --- a/python/packages/autogen-core/src/autogen_core/_agent.py +++ b/python/packages/autogen-core/src/autogen_core/_agent.py @@ -1,9 +1,13 @@ -from typing import Any, Mapping, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Mapping, Protocol, runtime_checkable from ._agent_id import AgentId from ._agent_metadata import AgentMetadata from ._message_context import MessageContext +# Forward declaration for type checking only +if TYPE_CHECKING: + from ._agent_runtime import AgentRuntime + @runtime_checkable class Agent(Protocol): @@ -17,6 +21,15 @@ def id(self) -> AgentId: """ID of the agent.""" ... + async def bind_id_and_runtime(self, id: AgentId, runtime: "AgentRuntime") -> None: + """Function used to bind an Agent instance to an `AgentRuntime`. + + Args: + agent_id (AgentId): ID of the agent. + runtime (AgentRuntime): AgentRuntime instance to bind the agent to. + """ + ... + async def on_message(self, message: Any, ctx: MessageContext) -> Any: """Message handler for the agent. This should only be called by the runtime, not by other agents. diff --git a/python/packages/autogen-core/src/autogen_core/_agent_instantiation.py b/python/packages/autogen-core/src/autogen_core/_agent_instantiation.py index 71921225cfbd..a8904a42da56 100644 --- a/python/packages/autogen-core/src/autogen_core/_agent_instantiation.py +++ b/python/packages/autogen-core/src/autogen_core/_agent_instantiation.py @@ -118,3 +118,9 @@ def current_agent_id(cls) -> AgentId: raise RuntimeError( "AgentInstantiationContext.agent_id() must be called within an instantiation context such as when the AgentRuntime is instantiating an agent. Mostly likely this was caused by directly instantiating an agent instead of using the AgentRuntime to do so." ) from e + + @classmethod + def is_in_factory_call(cls) -> bool: + if cls._AGENT_INSTANTIATION_CONTEXT_VAR.get(None) is None: + return False + return True diff --git a/python/packages/autogen-core/src/autogen_core/_agent_runtime.py b/python/packages/autogen-core/src/autogen_core/_agent_runtime.py index 6510d84fbe17..d4bac4a9c0ab 100644 --- a/python/packages/autogen-core/src/autogen_core/_agent_runtime.py +++ b/python/packages/autogen-core/src/autogen_core/_agent_runtime.py @@ -62,7 +62,7 @@ async def publish_message( Args: message (Any): The message to publish. - topic (TopicId): The topic to publish the message to. + topic_id (TopicId): The topic to publish the message to. sender (AgentId | None, optional): The agent which sent the message. Defaults to None. cancellation_token (CancellationToken | None, optional): Token used to cancel an in progress. Defaults to None. message_id (str | None, optional): The message id. If None, a new message id will be generated. Defaults to None. This message id must be unique. and is recommended to be a UUID. @@ -130,6 +130,60 @@ async def main() -> None: """ ... + async def register_agent_instance( + self, + agent_instance: Agent, + agent_id: AgentId, + ) -> AgentId: + """Register an agent instance with the runtime. The type may be reused, but each agent_id must be unique. All agent instances within a type must be of the same object type. This API does not add any subscriptions. + + .. note:: + + This is a low level API and usually the agent class's `register_instance` method should be used instead, as this also handles subscriptions automatically. + + Example: + + .. code-block:: python + + from dataclasses import dataclass + + from autogen_core import AgentId, AgentRuntime, MessageContext, RoutedAgent, event + from autogen_core.models import UserMessage + + + @dataclass + class MyMessage: + content: str + + + class MyAgent(RoutedAgent): + def __init__(self) -> None: + super().__init__("My core agent") + + @event + async def handler(self, message: UserMessage, context: MessageContext) -> None: + print("Event received: ", message.content) + + + async def main() -> None: + runtime: AgentRuntime = ... # type: ignore + agent = MyAgent() + await runtime.register_agent_instance( + agent_instance=agent, agent_id=AgentId(type="my_agent", key="default") + ) + + + import asyncio + + asyncio.run(main()) + + + Args: + agent_instance (Agent): A concrete instance of the agent. + agent_id (AgentId): The agent's identifier. The agent's type is `agent_id.type`. + """ + ... + # TODO: uncomment out the following type ignore when this is fixed in mypy: https://github.com/python/mypy/issues/3737 async def try_get_underlying_agent_instance(self, id: AgentId, type: Type[T] = Agent) -> T: # type: ignore[assignment] """Try to get the underlying agent instance by name and namespace. This is generally discouraged (hence the long name), but can be useful in some cases. diff --git a/python/packages/autogen-core/src/autogen_core/_base_agent.py b/python/packages/autogen-core/src/autogen_core/_base_agent.py index bffb61b876bb..0ad0bc60776c 100644 --- a/python/packages/autogen-core/src/autogen_core/_base_agent.py +++ b/python/packages/autogen-core/src/autogen_core/_base_agent.py @@ -21,6 +21,7 @@ from ._subscription_context import SubscriptionInstantiationContext from ._topic import TopicId from ._type_prefix_subscription import TypePrefixSubscription +from ._type_subscription import TypeSubscription T = TypeVar("T", bound=Agent) @@ -82,20 +83,25 @@ def metadata(self) -> AgentMetadata: return AgentMetadata(key=self._id.key, type=self._id.type, description=self._description) def __init__(self, description: str) -> None: - try: - runtime = AgentInstantiationContext.current_runtime() - id = AgentInstantiationContext.current_agent_id() - except LookupError as e: - raise RuntimeError( - "BaseAgent must be instantiated within the context of an AgentRuntime. It cannot be directly instantiated." - ) from e - - self._runtime: AgentRuntime = runtime - self._id: AgentId = id + if AgentInstantiationContext.is_in_factory_call(): + self._runtime: AgentRuntime = AgentInstantiationContext.current_runtime() + self._id = AgentInstantiationContext.current_agent_id() if not isinstance(description, str): raise ValueError("Agent description must be a string") self._description = description + async def bind_id_and_runtime(self, id: AgentId, runtime: AgentRuntime) -> None: + if hasattr(self, "_id"): + if self._id != id: + raise RuntimeError("Agent is already bound to a different ID") + + if hasattr(self, "_runtime"): + if self._runtime != runtime: + raise RuntimeError("Agent is already bound to a different runtime") + + self._id = id + self._runtime = runtime + @property def type(self) -> str: return self.id.type @@ -155,6 +161,56 @@ async def load_state(self, state: Mapping[str, Any]) -> None: async def close(self) -> None: pass + async def register_instance( + self, + runtime: AgentRuntime, + agent_id: AgentId, + *, + skip_class_subscriptions: bool = True, + skip_direct_message_subscription: bool = False, + ) -> AgentId: + """ + This function is similar to `register` but is used for registering an instance of an agent. A subscription based on the agent ID is created and added to the runtime. + """ + agent_id = await runtime.register_agent_instance(agent_instance=self, agent_id=agent_id) + + id_subscription = TypeSubscription(topic_type=agent_id.key, agent_type=agent_id.type) + await runtime.add_subscription(id_subscription) + + if not skip_class_subscriptions: + with SubscriptionInstantiationContext.populate_context(AgentType(agent_id.type)): + subscriptions: List[Subscription] = [] + for unbound_subscription in self._unbound_subscriptions(): + subscriptions_list_result = unbound_subscription() + if inspect.isawaitable(subscriptions_list_result): + subscriptions_list = await subscriptions_list_result + else: + subscriptions_list = subscriptions_list_result + + subscriptions.extend(subscriptions_list) + for subscription in subscriptions: + await runtime.add_subscription(subscription) + + if not skip_direct_message_subscription: + # Additionally adds a special prefix subscription for this agent to receive direct messages + try: + await runtime.add_subscription( + TypePrefixSubscription( + # The prefix MUST include ":" to avoid collisions with other agents + topic_type_prefix=agent_id.type + ":", + agent_type=agent_id.type, + ) + ) + except ValueError: + # We don't care if the subscription already exists + pass + + # TODO: deduplication + for _message_type, serializer in self._handles_types(): + runtime.add_message_serializer(serializer) + + return agent_id + @classmethod async def register( cls, diff --git a/python/packages/autogen-core/src/autogen_core/_component_config.py b/python/packages/autogen-core/src/autogen_core/_component_config.py index ce768ddafa68..bb603a839ed6 100644 --- a/python/packages/autogen-core/src/autogen_core/_component_config.py +++ b/python/packages/autogen-core/src/autogen_core/_component_config.py @@ -7,7 +7,7 @@ from pydantic import BaseModel from typing_extensions import Self, TypeVar -ComponentType = Literal["model", "agent", "tool", "termination", "token_provider"] | str +ComponentType = Literal["model", "agent", "tool", "termination", "token_provider", "workbench"] | str ConfigT = TypeVar("ConfigT", bound=BaseModel) FromConfigT = TypeVar("FromConfigT", bound=BaseModel, contravariant=True) ToConfigT = TypeVar("ToConfigT", bound=BaseModel, covariant=True) diff --git a/python/packages/autogen-core/src/autogen_core/_constants.py b/python/packages/autogen-core/src/autogen_core/_constants.py index 8fc4580c051d..06f3ab01c430 100644 --- a/python/packages/autogen-core/src/autogen_core/_constants.py +++ b/python/packages/autogen-core/src/autogen_core/_constants.py @@ -1,5 +1,5 @@ ROOT_LOGGER_NAME = "autogen_core" -"""str: Logger name used for structured event logging""" +"""str: Logger name used for root logger""" EVENT_LOGGER_NAME = "autogen_core.events" """str: Logger name used for structured event logging""" diff --git a/python/packages/autogen-core/src/autogen_core/_function_utils.py b/python/packages/autogen-core/src/autogen_core/_function_utils.py index cbf157d97a08..891027842794 100644 --- a/python/packages/autogen-core/src/autogen_core/_function_utils.py +++ b/python/packages/autogen-core/src/autogen_core/_function_utils.py @@ -155,7 +155,7 @@ def get_required_params(typed_signature: inspect.Signature) -> List[str]: """Get the required parameters of a function Args: - signature: The signature of the function as returned by inspect.signature + typed_signature: The signature of the function as returned by inspect.signature Returns: A list of the required parameters of the function @@ -167,7 +167,7 @@ def get_default_values(typed_signature: inspect.Signature) -> Dict[str, Any]: """Get default values of parameters of a function Args: - signature: The signature of the function as returned by inspect.signature + typed_signature: The signature of the function as returned by inspect.signature Returns: A dictionary of the default values of the parameters of the function @@ -184,7 +184,8 @@ def get_parameters( Args: required: The required parameters of the function - hints: The type hints of the function as returned by typing.get_type_hints + param_annotations: A dictionary of the type annotations of the parameters of the function + default_values: The default values of the parameters of the function Returns: A Pydantic model for the parameters of the function diff --git a/python/packages/autogen-core/src/autogen_core/_routed_agent.py b/python/packages/autogen-core/src/autogen_core/_routed_agent.py index a5908278cab9..cc4c114909aa 100644 --- a/python/packages/autogen-core/src/autogen_core/_routed_agent.py +++ b/python/packages/autogen-core/src/autogen_core/_routed_agent.py @@ -123,7 +123,7 @@ def decorator( raise AssertionError("message parameter not found in function signature") if "return" not in type_hints: - raise AssertionError("return not found in function signature") + raise AssertionError("return parameter not found in function signature") # Get the type of the message parameter target_types = get_types(type_hints["message"]) @@ -243,7 +243,7 @@ def decorator( raise AssertionError("message parameter not found in function signature") if "return" not in type_hints: - raise AssertionError("return not found in function signature") + raise AssertionError("return parameter not found in function signature") # Get the type of the message parameter target_types = get_types(type_hints["message"]) @@ -363,7 +363,7 @@ def decorator( raise AssertionError("message parameter not found in function signature") if "return" not in type_hints: - raise AssertionError("return not found in function signature") + raise AssertionError("return parameter not found in function signature") # Get the type of the message parameter target_types = get_types(type_hints["message"]) diff --git a/python/packages/autogen-core/src/autogen_core/_single_threaded_agent_runtime.py b/python/packages/autogen-core/src/autogen_core/_single_threaded_agent_runtime.py index 9610e7f54ebc..3a8a8d714ff0 100644 --- a/python/packages/autogen-core/src/autogen_core/_single_threaded_agent_runtime.py +++ b/python/packages/autogen-core/src/autogen_core/_single_threaded_agent_runtime.py @@ -2,6 +2,7 @@ import asyncio import inspect +import json import logging import sys import uuid @@ -159,6 +160,7 @@ class SingleThreadedAgentRuntime(AgentRuntime): intervention_handlers (List[InterventionHandler], optional): A list of intervention handlers that can intercept messages before they are sent or published. Defaults to None. tracer_provider (TracerProvider, optional): The tracer provider to use for tracing. Defaults to None. + Additionally, you can set environment variable `AUTOGEN_DISABLE_RUNTIME_TRACING` to `true` to disable the agent runtime telemetry if you don't have access to the runtime constructor. For example, if you are using `ComponentConfig`. ignore_unhandled_exceptions (bool, optional): Whether to ignore unhandled exceptions in that occur in agent event handlers. Any background exceptions will be raised on the next call to `process_next` or from an awaited `stop`, `stop_when_idle` or `stop_when`. Note, this does not apply to RPC handlers. Defaults to True. Examples: @@ -265,6 +267,7 @@ def __init__( self._serialization_registry = SerializationRegistry() self._ignore_unhandled_handler_exceptions = ignore_unhandled_exceptions self._background_exception: BaseException | None = None + self._agent_instance_types: Dict[str, Type[Agent]] = {} @property def unprocessed_messages_count( @@ -276,6 +279,55 @@ def unprocessed_messages_count( def _known_agent_names(self) -> Set[str]: return set(self._agent_factories.keys()) + async def _create_otel_attributes( + self, + sender_agent_id: AgentId | None = None, + recipient_agent_id: AgentId | None = None, + message_context: MessageContext | None = None, + message: Any = None, + ) -> Mapping[str, str]: + """Create OpenTelemetry attributes for the given agent and message. + + Args: + sender_agent (Agent, optional): The sender agent instance. + recipient_agent (Agent, optional): The recipient agent instance. + message (Any): The message instance. + + Returns: + Attributes: A dictionary of OpenTelemetry attributes. + """ + if not sender_agent_id and not recipient_agent_id and not message: + return {} + attributes: Dict[str, str] = {} + if sender_agent_id: + sender_agent = await self._get_agent(sender_agent_id) + attributes["sender_agent_type"] = sender_agent.id.type + attributes["sender_agent_class"] = sender_agent.__class__.__name__ + if recipient_agent_id: + recipient_agent = await self._get_agent(recipient_agent_id) + attributes["recipient_agent_type"] = recipient_agent.id.type + attributes["recipient_agent_class"] = recipient_agent.__class__.__name__ + + if message_context: + serialized_message_context = { + "sender": str(message_context.sender), + "topic_id": str(message_context.topic_id), + "is_rpc": message_context.is_rpc, + "message_id": message_context.message_id, + } + attributes["message_context"] = json.dumps(serialized_message_context) + + if message: + try: + serialized_message = self._try_serialize(message) + except Exception as e: + serialized_message = str(e) + else: + serialized_message = "No Message" + attributes["message"] = serialized_message + + return attributes + # Returns the response of the message async def send_message( self, @@ -311,6 +363,7 @@ async def send_message( future = asyncio.get_event_loop().create_future() if recipient.type not in self._known_agent_names: future.set_exception(Exception("Recipient not found")) + return await future content = message.__dict__ if hasattr(message, "__dict__") else message logger.info(f"Sending message of type {type(message).__name__} to {recipient.type}: {content}") @@ -440,7 +493,17 @@ async def _process_send(self, message_envelope: SendMessageEnvelope) -> None: cancellation_token=message_envelope.cancellation_token, message_id=message_envelope.message_id, ) - with self._tracer_helper.trace_block("process", recipient_agent.id, parent=message_envelope.metadata): + with self._tracer_helper.trace_block( + "process", + recipient_agent.id, + parent=message_envelope.metadata, + attributes=await self._create_otel_attributes( + sender_agent_id=message_envelope.sender, + recipient_agent_id=recipient, + message_context=message_context, + message=message_envelope.message, + ), + ): with MessageHandlerContext.populate_context(recipient_agent.id): response = await recipient_agent.on_message( message_envelope.message, @@ -527,7 +590,17 @@ async def _process_publish(self, message_envelope: PublishMessageEnvelope) -> No agent = await self._get_agent(agent_id) async def _on_message(agent: Agent, message_context: MessageContext) -> Any: - with self._tracer_helper.trace_block("process", agent.id, parent=message_envelope.metadata): + with self._tracer_helper.trace_block( + "process", + agent.id, + parent=message_envelope.metadata, + attributes=await self._create_otel_attributes( + sender_agent_id=message_envelope.sender, + recipient_agent_id=agent.id, + message_context=message_context, + message=message_envelope.message, + ), + ): with MessageHandlerContext.populate_context(agent.id): try: return await agent.on_message( @@ -557,7 +630,16 @@ async def _on_message(agent: Agent, message_context: MessageContext) -> Any: # TODO if responses are given for a publish async def _process_response(self, message_envelope: ResponseMessageEnvelope) -> None: - with self._tracer_helper.trace_block("ack", message_envelope.recipient, parent=message_envelope.metadata): + with self._tracer_helper.trace_block( + "ack", + message_envelope.recipient, + parent=message_envelope.metadata, + attributes=await self._create_otel_attributes( + sender_agent_id=message_envelope.sender, + recipient_agent_id=message_envelope.recipient, + message=message_envelope.message, + ), + ): content = ( message_envelope.message.__dict__ if hasattr(message_envelope.message, "__dict__") @@ -821,15 +903,42 @@ async def factory_wrapper() -> T: else: agent_instance = maybe_agent_instance - if expected_class is not None and type_func_alias(agent_instance) != expected_class: - raise ValueError("Factory registered using the wrong type.") - + if expected_class is not None and not issubclass(type_func_alias(agent_instance), expected_class): + raise ValueError( + f"Factory registered using the wrong type: expected {expected_class.__name__}, got {type_func_alias(agent_instance).__name__}" + ) return agent_instance self._agent_factories[type.type] = factory_wrapper return type + async def register_agent_instance( + self, + agent_instance: Agent, + agent_id: AgentId, + ) -> AgentId: + def agent_factory() -> Agent: + raise RuntimeError( + "Agent factory was invoked for an agent instance that was not registered. This is likely due to the agent type being incorrectly subscribed to a topic. If this exception occurs when publishing a message to the DefaultTopicId, then it is likely that `skip_class_subscriptions` needs to be turned off when registering the agent." + ) + + if agent_id in self._instantiated_agents: + raise ValueError(f"Agent with id {agent_id} already exists.") + + if agent_id.type not in self._agent_factories: + self._agent_factories[agent_id.type] = agent_factory + self._agent_instance_types[agent_id.type] = type_func_alias(agent_instance) + else: + if self._agent_factories[agent_id.type].__code__ != agent_factory.__code__: + raise ValueError("Agent factories and agent instances cannot be registered to the same type.") + if self._agent_instance_types[agent_id.type] != type_func_alias(agent_instance): + raise ValueError("Agent instances must be the same object type.") + + await agent_instance.bind_id_and_runtime(id=agent_id, runtime=self) + self._instantiated_agents[agent_id] = agent_instance + return agent_id + async def _invoke_agent_factory( self, agent_factory: Callable[[], T | Awaitable[T]] | Callable[[AgentRuntime, AgentId], T | Awaitable[T]], @@ -851,8 +960,7 @@ async def _invoke_agent_factory( raise ValueError("Agent factory must take 0 or 2 arguments.") if inspect.isawaitable(agent): - return cast(T, await agent) - + agent = cast(T, await agent) return agent except BaseException as e: diff --git a/python/packages/autogen-core/src/autogen_core/_telemetry/__init__.py b/python/packages/autogen-core/src/autogen_core/_telemetry/__init__.py index 6432edf43e2f..c67591a679e4 100644 --- a/python/packages/autogen-core/src/autogen_core/_telemetry/__init__.py +++ b/python/packages/autogen-core/src/autogen_core/_telemetry/__init__.py @@ -1,3 +1,8 @@ +from ._genai import ( + trace_create_agent_span, + trace_invoke_agent_span, + trace_tool_span, +) from ._propagation import ( EnvelopeMetadata, TelemetryMetadataContainer, @@ -14,4 +19,7 @@ "TelemetryMetadataContainer", "TraceHelper", "MessageRuntimeTracingConfig", + "trace_create_agent_span", + "trace_invoke_agent_span", + "trace_tool_span", ] diff --git a/python/packages/autogen-core/src/autogen_core/_telemetry/_genai.py b/python/packages/autogen-core/src/autogen_core/_telemetry/_genai.py new file mode 100644 index 000000000000..ccbb5a353f4c --- /dev/null +++ b/python/packages/autogen-core/src/autogen_core/_telemetry/_genai.py @@ -0,0 +1,214 @@ +from collections.abc import Generator +from contextlib import contextmanager +from enum import Enum +from typing import Any, Optional + +from opentelemetry import trace +from opentelemetry.trace import Span, SpanKind + +from .._agent_instantiation import AgentInstantiationContext + +# OpenTelemetry semantic convention constants for GenAI operations +# Copied from opentelemetry-semantic-conventions to avoid dependency + +# GenAI Agent attributes +GEN_AI_AGENT_DESCRIPTION = "gen_ai.agent.description" +GEN_AI_AGENT_ID = "gen_ai.agent.id" +GEN_AI_AGENT_NAME = "gen_ai.agent.name" + +# GenAI Operation attributes +GEN_AI_OPERATION_NAME = "gen_ai.operation.name" +GEN_AI_SYSTEM = "gen_ai.system" + +# GenAI Tool attributes +GEN_AI_TOOL_CALL_ID = "gen_ai.tool.call.id" +GEN_AI_TOOL_DESCRIPTION = "gen_ai.tool.description" +GEN_AI_TOOL_NAME = "gen_ai.tool.name" + +# Error attributes +ERROR_TYPE = "error.type" + + +class GenAiOperationNameValues(Enum): + """Enum for GenAI operation name values.""" + + CHAT = "chat" + CREATE_AGENT = "create_agent" + EMBEDDINGS = "embeddings" + EXECUTE_TOOL = "execute_tool" + GENERATE_CONTENT = "generate_content" + INVOKE_AGENT = "invoke_agent" + TEXT_COMPLETION = "text_completion" + + +# Constant for system name +GENAI_SYSTEM_AUTOGEN = "autogen" + + +@contextmanager +def trace_tool_span( + tool_name: str, + *, + tracer: Optional[trace.Tracer] = None, + parent: Optional[Span] = None, + tool_description: Optional[str] = None, + tool_call_id: Optional[str] = None, +) -> Generator[Span, Any, None]: + """Context manager to create a span for tool execution following the + OpenTelemetry Semantic conventions for generative AI systems. + + See the GenAI semantic conventions documentation: + `OpenTelemetry GenAI Semantic Conventions `__ + + .. warning:: + + The GenAI Semantic Conventions are still in incubation and + subject to changes in future releases. + + + Args: + tool_name (str): The name of the tool being executed. + tracer (Optional[trace.Tracer]): The tracer to use for creating the span. + parent (Optional[Span]): The parent span to link this span to. + tool_description (Optional[str]): A description of the tool. + tool_call_id (Optional[str]): A unique identifier for the tool call. + """ + if tracer is None: + tracer = trace.get_tracer("autogen-core") + span_attributes = { + GEN_AI_OPERATION_NAME: GenAiOperationNameValues.EXECUTE_TOOL.value, + GEN_AI_SYSTEM: GENAI_SYSTEM_AUTOGEN, + GEN_AI_TOOL_NAME: tool_name, + } + if tool_description is not None: + span_attributes[GEN_AI_TOOL_DESCRIPTION] = tool_description + if tool_call_id is not None: + span_attributes[GEN_AI_TOOL_CALL_ID] = tool_call_id + with tracer.start_as_current_span( + f"{GenAiOperationNameValues.EXECUTE_TOOL.value} {tool_name}", + kind=SpanKind.INTERNAL, + context=trace.set_span_in_context(parent) if parent else None, + attributes=span_attributes, + ) as span: + try: + yield span + except Exception as e: + # Set the exception details on the span if an error occurs + span.record_exception(e) + span.set_status(trace.Status(trace.StatusCode.ERROR, str(e))) + span.set_attribute(ERROR_TYPE, type(e).__name__) + raise + + +@contextmanager +def trace_create_agent_span( + agent_name: str, + *, + tracer: Optional[trace.Tracer] = None, + parent: Optional[Span] = None, + agent_id: Optional[str] = None, + agent_description: Optional[str] = None, +) -> Generator[Span, Any, None]: + """Context manager to create a span for agent creation following the + OpenTelemetry Semantic conventions for generative AI systems. + + See the GenAI semantic conventions documentation: + `OpenTelemetry GenAI Semantic Conventions `__ + + .. warning:: + + The GenAI Semantic Conventions are still in incubation and + subject to changes in future releases. + + Args: + agent_name (str): The name of the agent being created. + tracer (Optional[trace.Tracer]): The tracer to use for creating the span. + parent (Optional[Span]): The parent span to link this span to. + agent_id (Optional[str]): The unique identifier for the agent. + agent_description (Optional[str]): A description of the agent. + """ + if tracer is None: + tracer = trace.get_tracer("autogen-core") + span_attributes = { + GEN_AI_OPERATION_NAME: GenAiOperationNameValues.CREATE_AGENT.value, + GEN_AI_SYSTEM: GENAI_SYSTEM_AUTOGEN, + GEN_AI_AGENT_NAME: agent_name, + } + if agent_id is None: + # Try to see if we can get the agent ID from the current context + try: + agent_id = str(AgentInstantiationContext.current_agent_id()) + except RuntimeError: + agent_id = None + if agent_id is not None: + span_attributes[GEN_AI_AGENT_ID] = agent_id + if agent_description is not None: + span_attributes[GEN_AI_AGENT_DESCRIPTION] = agent_description + with tracer.start_as_current_span( + f"{GenAiOperationNameValues.CREATE_AGENT.value} {agent_name}", + kind=SpanKind.CLIENT, + context=trace.set_span_in_context(parent) if parent else None, + attributes=span_attributes, + ) as span: + try: + yield span + except Exception as e: + # Set the exception details on the span if an error occurs + span.record_exception(e) + span.set_status(trace.Status(trace.StatusCode.ERROR, str(e))) + span.set_attribute(ERROR_TYPE, type(e).__name__) + raise + + +@contextmanager +def trace_invoke_agent_span( + agent_name: str, + *, + tracer: Optional[trace.Tracer] = None, + parent: Optional[Span] = None, + agent_id: Optional[str] = None, + agent_description: Optional[str] = None, +) -> Generator[Span, Any, None]: + """Context manager to create a span for invoking an agent following the + OpenTelemetry Semantic conventions for generative AI systems. + + See the GenAI semantic conventions documentation: + `OpenTelemetry GenAI Semantic Conventions `__ + + .. warning:: + + The GenAI Semantic Conventions are still in incubation and + subject to changes in future releases. + + Args: + agent_name (str): The name of the agent being invoked. + tracer (Optional[trace.Tracer]): The tracer to use for creating the span. + parent (Optional[Span]): The parent span to link this span to. + agent_id (Optional[str]): The unique identifier for the agent. + agent_description (Optional[str]): A description of the agent. + """ + if tracer is None: + tracer = trace.get_tracer("autogen-core") + span_attributes = { + GEN_AI_OPERATION_NAME: GenAiOperationNameValues.INVOKE_AGENT.value, + GEN_AI_SYSTEM: GENAI_SYSTEM_AUTOGEN, + GEN_AI_AGENT_NAME: agent_name, + } + if agent_id is not None: + span_attributes[GEN_AI_AGENT_ID] = agent_id + if agent_description is not None: + span_attributes[GEN_AI_AGENT_DESCRIPTION] = agent_description + with tracer.start_as_current_span( + f"{GenAiOperationNameValues.INVOKE_AGENT.value} {agent_name}", + kind=SpanKind.CLIENT, + context=trace.set_span_in_context(parent) if parent else None, + attributes=span_attributes, + ) as span: + try: + yield span + except Exception as e: + # Set the exception details on the span if an error occurs + span.record_exception(e) + span.set_status(trace.Status(trace.StatusCode.ERROR, str(e))) + span.set_attribute(ERROR_TYPE, type(e).__name__) + raise diff --git a/python/packages/autogen-core/src/autogen_core/_telemetry/_tracing.py b/python/packages/autogen-core/src/autogen_core/_telemetry/_tracing.py index aa1835090aa1..2f73e6e5af0f 100644 --- a/python/packages/autogen-core/src/autogen_core/_telemetry/_tracing.py +++ b/python/packages/autogen-core/src/autogen_core/_telemetry/_tracing.py @@ -1,4 +1,5 @@ import contextlib +import os from typing import Dict, Generic, Iterator, Optional from opentelemetry.trace import NoOpTracerProvider, Span, SpanKind, TracerProvider, get_tracer_provider @@ -22,11 +23,18 @@ def __init__( tracer_provider: TracerProvider | None, instrumentation_builder_config: TracingConfig[Operation, Destination, ExtraAttributes], ) -> None: + self.instrumentation_builder_config = instrumentation_builder_config + + disable_runtime_tracing = os.environ.get("AUTOGEN_DISABLE_RUNTIME_TRACING") == "true" + if disable_runtime_tracing: + self.tracer_provider: TracerProvider = NoOpTracerProvider() + self.tracer = self.tracer_provider.get_tracer(f"autogen {instrumentation_builder_config.name}") + return + # Evaluate in order: first try tracer_provider param, then get_tracer_provider(), finally fallback to NoOp # This allows for nested tracing with a default tracer provided by the user self.tracer_provider = tracer_provider or get_tracer_provider() or NoOpTracerProvider() self.tracer = self.tracer_provider.get_tracer(f"autogen {instrumentation_builder_config.name}") - self.instrumentation_builder_config = instrumentation_builder_config @contextlib.contextmanager def trace_block( diff --git a/python/packages/autogen-core/src/autogen_core/logging.py b/python/packages/autogen-core/src/autogen_core/logging.py index 18811f8d2b5b..3f371a6f3bca 100644 --- a/python/packages/autogen-core/src/autogen_core/logging.py +++ b/python/packages/autogen-core/src/autogen_core/logging.py @@ -113,7 +113,7 @@ def __init__( completion_tokens: int, **kwargs: Any, ) -> None: - """To be used by model to log the call to the LLM. + """To be used by model clients to log the end of a stream. Args: response (Dict[str, Any]): The response of the call. Must be json serializable. diff --git a/python/packages/autogen-core/src/autogen_core/memory/_list_memory.py b/python/packages/autogen-core/src/autogen_core/memory/_list_memory.py index 5ad086550081..9dd1182c74c7 100644 --- a/python/packages/autogen-core/src/autogen_core/memory/_list_memory.py +++ b/python/packages/autogen-core/src/autogen_core/memory/_list_memory.py @@ -1,6 +1,6 @@ from typing import Any, List -from pydantic import BaseModel +from pydantic import BaseModel, Field from typing_extensions import Self from .._cancellation_token import CancellationToken @@ -15,7 +15,7 @@ class ListMemoryConfig(BaseModel): name: str | None = None """Optional identifier for this memory instance.""" - memory_contents: List[MemoryContent] = [] + memory_contents: List[MemoryContent] = Field(default_factory=list) """List of memory contents stored in this memory instance.""" diff --git a/python/packages/autogen-core/src/autogen_core/model_context/__init__.py b/python/packages/autogen-core/src/autogen_core/model_context/__init__.py index 0c8c7af5cf08..b6898614ec37 100644 --- a/python/packages/autogen-core/src/autogen_core/model_context/__init__.py +++ b/python/packages/autogen-core/src/autogen_core/model_context/__init__.py @@ -1,6 +1,7 @@ from ._buffered_chat_completion_context import BufferedChatCompletionContext from ._chat_completion_context import ChatCompletionContext, ChatCompletionContextState from ._head_and_tail_chat_completion_context import HeadAndTailChatCompletionContext +from ._token_limited_chat_completion_context import TokenLimitedChatCompletionContext from ._unbounded_chat_completion_context import ( UnboundedChatCompletionContext, ) @@ -10,5 +11,6 @@ "ChatCompletionContextState", "UnboundedChatCompletionContext", "BufferedChatCompletionContext", + "TokenLimitedChatCompletionContext", "HeadAndTailChatCompletionContext", ] diff --git a/python/packages/autogen-core/src/autogen_core/model_context/_buffered_chat_completion_context.py b/python/packages/autogen-core/src/autogen_core/model_context/_buffered_chat_completion_context.py index dcece60b1cd7..5d23f818a56c 100644 --- a/python/packages/autogen-core/src/autogen_core/model_context/_buffered_chat_completion_context.py +++ b/python/packages/autogen-core/src/autogen_core/model_context/_buffered_chat_completion_context.py @@ -41,7 +41,9 @@ async def get_messages(self) -> List[LLMMessage]: return messages def _to_config(self) -> BufferedChatCompletionContextConfig: - return BufferedChatCompletionContextConfig(buffer_size=self._buffer_size, initial_messages=self._messages) + return BufferedChatCompletionContextConfig( + buffer_size=self._buffer_size, initial_messages=self._initial_messages + ) @classmethod def _from_config(cls, config: BufferedChatCompletionContextConfig) -> Self: diff --git a/python/packages/autogen-core/src/autogen_core/model_context/_chat_completion_context.py b/python/packages/autogen-core/src/autogen_core/model_context/_chat_completion_context.py index 1b14b8e85f68..84871f1548a9 100644 --- a/python/packages/autogen-core/src/autogen_core/model_context/_chat_completion_context.py +++ b/python/packages/autogen-core/src/autogen_core/model_context/_chat_completion_context.py @@ -47,7 +47,10 @@ async def get_messages(self) -> List[LLMMessage]: component_type = "chat_completion_context" def __init__(self, initial_messages: List[LLMMessage] | None = None) -> None: - self._messages: List[LLMMessage] = initial_messages or [] + self._messages: List[LLMMessage] = [] + if initial_messages is not None: + self._messages.extend(initial_messages) + self._initial_messages = initial_messages async def add_message(self, message: LLMMessage) -> None: """Add a message to the context.""" diff --git a/python/packages/autogen-core/src/autogen_core/model_context/_head_and_tail_chat_completion_context.py b/python/packages/autogen-core/src/autogen_core/model_context/_head_and_tail_chat_completion_context.py index a37d5927b19f..75493618e90d 100644 --- a/python/packages/autogen-core/src/autogen_core/model_context/_head_and_tail_chat_completion_context.py +++ b/python/packages/autogen-core/src/autogen_core/model_context/_head_and_tail_chat_completion_context.py @@ -68,7 +68,7 @@ async def get_messages(self) -> List[LLMMessage]: def _to_config(self) -> HeadAndTailChatCompletionContextConfig: return HeadAndTailChatCompletionContextConfig( - head_size=self._head_size, tail_size=self._tail_size, initial_messages=self._messages + head_size=self._head_size, tail_size=self._tail_size, initial_messages=self._initial_messages ) @classmethod diff --git a/python/packages/autogen-core/src/autogen_core/model_context/_token_limited_chat_completion_context.py b/python/packages/autogen-core/src/autogen_core/model_context/_token_limited_chat_completion_context.py new file mode 100644 index 000000000000..b8a0258a4d62 --- /dev/null +++ b/python/packages/autogen-core/src/autogen_core/model_context/_token_limited_chat_completion_context.py @@ -0,0 +1,94 @@ +from typing import List + +from pydantic import BaseModel +from typing_extensions import Self + +from .._component_config import Component, ComponentModel +from ..models import ChatCompletionClient, FunctionExecutionResultMessage, LLMMessage +from ..tools import ToolSchema +from ._chat_completion_context import ChatCompletionContext + + +class TokenLimitedChatCompletionContextConfig(BaseModel): + model_client: ComponentModel + token_limit: int | None = None + tool_schema: List[ToolSchema] | None = None + initial_messages: List[LLMMessage] | None = None + + +class TokenLimitedChatCompletionContext(ChatCompletionContext, Component[TokenLimitedChatCompletionContextConfig]): + """(Experimental) A token based chat completion context maintains a view of the context up to a token limit. + + .. note:: + + Added in v0.4.10. This is an experimental component and may change in the future. + + Args: + model_client (ChatCompletionClient): The model client to use for token counting. + The model client must implement the :meth:`~autogen_core.models.ChatCompletionClient.count_tokens` + and :meth:`~autogen_core.models.ChatCompletionClient.remaining_tokens` methods. + token_limit (int | None): The maximum number of tokens to keep in the context + using the :meth:`~autogen_core.models.ChatCompletionClient.count_tokens` method. + If None, the context will be limited by the model client using the + :meth:`~autogen_core.models.ChatCompletionClient.remaining_tokens` method. + tools (List[ToolSchema] | None): A list of tool schema to use in the context. + initial_messages (List[LLMMessage] | None): A list of initial messages to include in the context. + + """ + + component_config_schema = TokenLimitedChatCompletionContextConfig + component_provider_override = "autogen_core.model_context.TokenLimitedChatCompletionContext" + + def __init__( + self, + model_client: ChatCompletionClient, + *, + token_limit: int | None = None, + tool_schema: List[ToolSchema] | None = None, + initial_messages: List[LLMMessage] | None = None, + ) -> None: + super().__init__(initial_messages) + if token_limit is not None and token_limit <= 0: + raise ValueError("token_limit must be greater than 0.") + self._token_limit = token_limit + self._model_client = model_client + self._tool_schema = tool_schema or [] + + async def get_messages(self) -> List[LLMMessage]: + """Get at most `token_limit` tokens in recent messages. If the token limit is not + provided, then return as many messages as the remaining token allowed by the model client.""" + messages = list(self._messages) + if self._token_limit is None: + remaining_tokens = self._model_client.remaining_tokens(messages, tools=self._tool_schema) + while remaining_tokens < 0 and len(messages) > 0: + middle_index = len(messages) // 2 + messages.pop(middle_index) + remaining_tokens = self._model_client.remaining_tokens(messages, tools=self._tool_schema) + else: + token_count = self._model_client.count_tokens(messages, tools=self._tool_schema) + while token_count > self._token_limit and len(messages) > 0: + middle_index = len(messages) // 2 + messages.pop(middle_index) + token_count = self._model_client.count_tokens(messages, tools=self._tool_schema) + if messages and isinstance(messages[0], FunctionExecutionResultMessage): + # Handle the first message is a function call result message. + # Remove the first message from the list. + messages = messages[1:] + return messages + + def _to_config(self) -> TokenLimitedChatCompletionContextConfig: + return TokenLimitedChatCompletionContextConfig( + model_client=self._model_client.dump_component(), + token_limit=self._token_limit, + tool_schema=self._tool_schema, + initial_messages=self._initial_messages, + ) + + @classmethod + def _from_config(cls, config: TokenLimitedChatCompletionContextConfig) -> Self: + return cls( + model_client=ChatCompletionClient.load_component(config.model_client), + token_limit=config.token_limit, + tool_schema=config.tool_schema, + initial_messages=config.initial_messages, + ) diff --git a/python/packages/autogen-core/src/autogen_core/model_context/_unbounded_chat_completion_context.py b/python/packages/autogen-core/src/autogen_core/model_context/_unbounded_chat_completion_context.py index 4bc26db46ae6..a2f409719f7c 100644 --- a/python/packages/autogen-core/src/autogen_core/model_context/_unbounded_chat_completion_context.py +++ b/python/packages/autogen-core/src/autogen_core/model_context/_unbounded_chat_completion_context.py @@ -9,7 +9,7 @@ class UnboundedChatCompletionContextConfig(BaseModel): - pass + initial_messages: List[LLMMessage] | None = None class UnboundedChatCompletionContext(ChatCompletionContext, Component[UnboundedChatCompletionContextConfig]): @@ -23,8 +23,8 @@ async def get_messages(self) -> List[LLMMessage]: return self._messages def _to_config(self) -> UnboundedChatCompletionContextConfig: - return UnboundedChatCompletionContextConfig() + return UnboundedChatCompletionContextConfig(initial_messages=self._initial_messages) @classmethod def _from_config(cls, config: UnboundedChatCompletionContextConfig) -> Self: - return cls() + return cls(initial_messages=config.initial_messages) diff --git a/python/packages/autogen-core/src/autogen_core/models/_model_client.py b/python/packages/autogen-core/src/autogen_core/models/_model_client.py index ab014ad45c93..4b328a301b42 100644 --- a/python/packages/autogen-core/src/autogen_core/models/_model_client.py +++ b/python/packages/autogen-core/src/autogen_core/models/_model_client.py @@ -18,38 +18,77 @@ class ModelFamily: This namespace class holds constants for the model families that AutoGen understands. Other families definitely exist and can be represented by a string, however, AutoGen will treat them as unknown.""" + GPT_41 = "gpt-41" + GPT_45 = "gpt-45" GPT_4O = "gpt-4o" O1 = "o1" O3 = "o3" + O4 = "o4" GPT_4 = "gpt-4" GPT_35 = "gpt-35" R1 = "r1" GEMINI_1_5_FLASH = "gemini-1.5-flash" GEMINI_1_5_PRO = "gemini-1.5-pro" GEMINI_2_0_FLASH = "gemini-2.0-flash" + GEMINI_2_5_PRO = "gemini-2.5-pro" + GEMINI_2_5_FLASH = "gemini-2.5-flash" CLAUDE_3_HAIKU = "claude-3-haiku" CLAUDE_3_SONNET = "claude-3-sonnet" CLAUDE_3_OPUS = "claude-3-opus" - CLAUDE_3_5_HAIKU = "claude-3.5-haiku" - CLAUDE_3_5_SONNET = "claude-3.5-sonnet" - CLAUDE_3_7_SONNET = "claude-3.7-sonnet" + CLAUDE_3_5_HAIKU = "claude-3-5-haiku" + CLAUDE_3_5_SONNET = "claude-3-5-sonnet" + CLAUDE_3_7_SONNET = "claude-3-7-sonnet" + CLAUDE_4_OPUS = "claude-4-opus" + CLAUDE_4_SONNET = "claude-4-sonnet" + LLAMA_3_3_8B = "llama-3.3-8b" + LLAMA_3_3_70B = "llama-3.3-70b" + LLAMA_4_SCOUT = "llama-4-scout" + LLAMA_4_MAVERICK = "llama-4-maverick" + CODESRAL = "codestral" + OPEN_CODESRAL_MAMBA = "open-codestral-mamba" + MISTRAL = "mistral" + MINISTRAL = "ministral" + PIXTRAL = "pixtral" UNKNOWN = "unknown" ANY: TypeAlias = Literal[ + # openai_models + "gpt-41", + "gpt-45", "gpt-4o", "o1", "o3", + "o4", "gpt-4", "gpt-35", "r1", + # google_models "gemini-1.5-flash", "gemini-1.5-pro", "gemini-2.0-flash", + "gemini-2.5-pro", + "gemini-2.5-flash", + # anthropic_models "claude-3-haiku", "claude-3-sonnet", "claude-3-opus", - "claude-3.5-haiku", - "claude-3.5-sonnet", + "claude-3-5-haiku", + "claude-3-5-sonnet", + "claude-3-7-sonnet", + "claude-4-opus", + "claude-4-sonnet", + # llama_models + "llama-3.3-8b", + "llama-3.3-70b", + "llama-4-scout", + "llama-4-maverick", + # mistral_models + "codestral", + "open-codestral-mamba", + "mistral", + "ministral", + "pixtral", + # unknown "unknown", ] @@ -64,6 +103,9 @@ def is_claude(family: str) -> bool: ModelFamily.CLAUDE_3_OPUS, ModelFamily.CLAUDE_3_5_HAIKU, ModelFamily.CLAUDE_3_5_SONNET, + ModelFamily.CLAUDE_3_7_SONNET, + ModelFamily.CLAUDE_4_OPUS, + ModelFamily.CLAUDE_4_SONNET, ) @staticmethod @@ -72,18 +114,42 @@ def is_gemini(family: str) -> bool: ModelFamily.GEMINI_1_5_FLASH, ModelFamily.GEMINI_1_5_PRO, ModelFamily.GEMINI_2_0_FLASH, + ModelFamily.GEMINI_2_5_PRO, + ModelFamily.GEMINI_2_5_FLASH, ) @staticmethod def is_openai(family: str) -> bool: return family in ( + ModelFamily.GPT_45, + ModelFamily.GPT_41, ModelFamily.GPT_4O, ModelFamily.O1, ModelFamily.O3, + ModelFamily.O4, ModelFamily.GPT_4, ModelFamily.GPT_35, ) + @staticmethod + def is_llama(family: str) -> bool: + return family in ( + ModelFamily.LLAMA_3_3_8B, + ModelFamily.LLAMA_3_3_70B, + ModelFamily.LLAMA_4_SCOUT, + ModelFamily.LLAMA_4_MAVERICK, + ) + + @staticmethod + def is_mistral(family: str) -> bool: + return family in ( + ModelFamily.CODESRAL, + ModelFamily.OPEN_CODESRAL_MAMBA, + ModelFamily.MISTRAL, + ModelFamily.MINISTRAL, + ModelFamily.PIXTRAL, + ) + @deprecated("Use the ModelInfo class instead ModelCapabilities.") class ModelCapabilities(TypedDict, total=False): @@ -109,6 +175,8 @@ class ModelInfo(TypedDict, total=False): """Model family should be one of the constants from :py:class:`ModelFamily` or a string representing an unknown model family.""" structured_output: Required[bool] """True if the model supports structured output, otherwise False. This is different to json_output.""" + multiple_system_messages: Optional[bool] + """True if the model supports multiple, non-consecutive system messages, otherwise False.""" def validate_model_info(model_info: ModelInfo) -> None: @@ -143,8 +211,7 @@ async def create( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], - # None means do not override the default - # A value means to override the client default - often specified in the constructor + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -154,6 +221,7 @@ async def create( Args: messages (Sequence[LLMMessage]): The messages to send to the model. tools (Sequence[Tool | ToolSchema], optional): The tools to use with the model. Defaults to []. + tool_choice (Tool | Literal["auto", "required", "none"], optional): A single Tool object to force the model to use, "auto" to let the model choose any available tool, "required" to force tool usage, or "none" to disable tool usage. Defaults to "auto". json_output (Optional[bool | type[BaseModel]], optional): Whether to use JSON mode, structured output, or neither. Defaults to None. If set to a `Pydantic BaseModel `_ type, it will be used as the output type for structured output. @@ -173,8 +241,7 @@ def create_stream( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], - # None means do not override the default - # A value means to override the client default - often specified in the constructor + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -184,6 +251,7 @@ def create_stream( Args: messages (Sequence[LLMMessage]): The messages to send to the model. tools (Sequence[Tool | ToolSchema], optional): The tools to use with the model. Defaults to []. + tool_choice (Tool | Literal["auto", "required", "none"], optional): A single Tool object to force the model to use, "auto" to let the model choose any available tool, "required" to force tool usage, or "none" to disable tool usage. Defaults to "auto". json_output (Optional[bool | type[BaseModel]], optional): Whether to use JSON mode, structured output, or neither. Defaults to None. If set to a `Pydantic BaseModel `_ type, it will be used as the output type for structured output. diff --git a/python/packages/autogen-core/src/autogen_core/tool_agent/_tool_agent.py b/python/packages/autogen-core/src/autogen_core/tool_agent/_tool_agent.py index 3aca56d6f49b..2ddb8dc2da4c 100644 --- a/python/packages/autogen-core/src/autogen_core/tool_agent/_tool_agent.py +++ b/python/packages/autogen-core/src/autogen_core/tool_agent/_tool_agent.py @@ -83,7 +83,9 @@ async def handle_function_call(self, message: FunctionCall, ctx: MessageContext) else: try: arguments = json.loads(message.arguments) - result = await tool.run_json(args=arguments, cancellation_token=ctx.cancellation_token) + result = await tool.run_json( + args=arguments, cancellation_token=ctx.cancellation_token, call_id=message.id + ) result_as_str = tool.return_value_as_string(result) except json.JSONDecodeError as e: raise InvalidToolArgumentsException( diff --git a/python/packages/autogen-core/src/autogen_core/tools/__init__.py b/python/packages/autogen-core/src/autogen_core/tools/__init__.py index 52a9d725f7d6..aee634e1fe24 100644 --- a/python/packages/autogen-core/src/autogen_core/tools/__init__.py +++ b/python/packages/autogen-core/src/autogen_core/tools/__init__.py @@ -1,11 +1,31 @@ -from ._base import BaseTool, BaseToolWithState, ParametersSchema, Tool, ToolSchema +from ._base import ( + BaseStreamTool, + BaseTool, + BaseToolWithState, + ParametersSchema, + StreamTool, + Tool, + ToolOverride, + ToolSchema, +) from ._function_tool import FunctionTool +from ._static_workbench import StaticStreamWorkbench, StaticWorkbench +from ._workbench import ImageResultContent, TextResultContent, ToolResult, Workbench __all__ = [ "Tool", + "StreamTool", "ToolSchema", "ParametersSchema", "BaseTool", "BaseToolWithState", + "BaseStreamTool", "FunctionTool", + "Workbench", + "ToolResult", + "TextResultContent", + "ImageResultContent", + "StaticWorkbench", + "StaticStreamWorkbench", + "ToolOverride", ] diff --git a/python/packages/autogen-core/src/autogen_core/tools/_base.py b/python/packages/autogen-core/src/autogen_core/tools/_base.py index 1843f246f203..d2ea76e21da1 100644 --- a/python/packages/autogen-core/src/autogen_core/tools/_base.py +++ b/python/packages/autogen-core/src/autogen_core/tools/_base.py @@ -2,16 +2,28 @@ import logging from abc import ABC, abstractmethod from collections.abc import Sequence -from typing import Any, Dict, Generic, Mapping, Protocol, Type, TypedDict, TypeVar, cast, runtime_checkable +from typing import ( + Any, + AsyncGenerator, + Dict, + Generic, + Mapping, + Optional, + Protocol, + Type, + TypeVar, + cast, + runtime_checkable, +) import jsonref -from opentelemetry.trace import get_tracer from pydantic import BaseModel -from typing_extensions import NotRequired +from typing_extensions import NotRequired, TypedDict from .. import EVENT_LOGGER_NAME, CancellationToken from .._component_config import ComponentBase from .._function_utils import normalize_annotated_type +from .._telemetry import trace_tool_span from ..logging import ToolCallEvent T = TypeVar("T", bound=BaseModel, contravariant=True) @@ -33,6 +45,13 @@ class ToolSchema(TypedDict): strict: NotRequired[bool] +class ToolOverride(BaseModel): + """Override configuration for a tool's name and/or description.""" + + name: Optional[str] = None + description: Optional[str] = None + + @runtime_checkable class Tool(Protocol): @property @@ -52,16 +71,26 @@ def state_type(self) -> Type[BaseModel] | None: ... def return_value_as_string(self, value: Any) -> str: ... - async def run_json(self, args: Mapping[str, Any], cancellation_token: CancellationToken) -> Any: ... + async def run_json( + self, args: Mapping[str, Any], cancellation_token: CancellationToken, call_id: str | None = None + ) -> Any: ... - def save_state_json(self) -> Mapping[str, Any]: ... + async def save_state_json(self) -> Mapping[str, Any]: ... - def load_state_json(self, state: Mapping[str, Any]) -> None: ... + async def load_state_json(self, state: Mapping[str, Any]) -> None: ... + + +@runtime_checkable +class StreamTool(Tool, Protocol): + def run_json_stream( + self, args: Mapping[str, Any], cancellation_token: CancellationToken, call_id: str | None = None + ) -> AsyncGenerator[Any, None]: ... ArgsT = TypeVar("ArgsT", bound=BaseModel, contravariant=True) ReturnT = TypeVar("ReturnT", bound=BaseModel, covariant=True) StateT = TypeVar("StateT", bound=BaseModel) +StreamT = TypeVar("StreamT", bound=BaseModel, covariant=True) class BaseTool(ABC, Tool, Generic[ArgsT, ReturnT], ComponentBase[BaseModel]): @@ -147,14 +176,23 @@ def return_value_as_string(self, value: Any) -> str: @abstractmethod async def run(self, args: ArgsT, cancellation_token: CancellationToken) -> ReturnT: ... - async def run_json(self, args: Mapping[str, Any], cancellation_token: CancellationToken) -> Any: - with get_tracer("base_tool").start_as_current_span( - self._name, - attributes={ - "tool_name": self._name, - "tool_description": self._description, - "tool_args": json.dumps(args), - }, + async def run_json( + self, args: Mapping[str, Any], cancellation_token: CancellationToken, call_id: str | None = None + ) -> Any: + """Run the tool with the provided arguments in a dictionary. + + Args: + args (Mapping[str, Any]): The arguments to pass to the tool. + cancellation_token (CancellationToken): A token to cancel the operation if needed. + call_id (str | None): An optional identifier for the tool call, used for tracing. + + Returns: + Any: The return value of the tool's run method. + """ + with trace_tool_span( + tool_name=self._name, + tool_description=self._description, + tool_call_id=call_id, ): # Execute the tool's run method return_value = await self.run(self._args_type.model_validate(args), cancellation_token) @@ -169,13 +207,66 @@ async def run_json(self, args: Mapping[str, Any], cancellation_token: Cancellati return return_value - def save_state_json(self) -> Mapping[str, Any]: + async def save_state_json(self) -> Mapping[str, Any]: return {} - def load_state_json(self, state: Mapping[str, Any]) -> None: + async def load_state_json(self, state: Mapping[str, Any]) -> None: pass +class BaseStreamTool( + BaseTool[ArgsT, ReturnT], StreamTool, ABC, Generic[ArgsT, StreamT, ReturnT], ComponentBase[BaseModel] +): + component_type = "tool" + + @abstractmethod + def run_stream(self, args: ArgsT, cancellation_token: CancellationToken) -> AsyncGenerator[StreamT | ReturnT, None]: + """Run the tool with the provided arguments and return a stream of data and end with the final return value.""" + ... + + async def run_json_stream( + self, + args: Mapping[str, Any], + cancellation_token: CancellationToken, + call_id: str | None = None, + ) -> AsyncGenerator[StreamT | ReturnT, None]: + """Run the tool with the provided arguments in a dictionary and return a stream of data + from the tool's :meth:`run_stream` method and end with the final return value. + + Args: + args (Mapping[str, Any]): The arguments to pass to the tool. + cancellation_token (CancellationToken): A token to cancel the operation if needed. + call_id (str | None): An optional identifier for the tool call, used for tracing. + + Returns: + AsyncGenerator[StreamT | ReturnT, None]: A generator yielding results from the tool's :meth:`run_stream` method. + """ + return_value: ReturnT | StreamT | None = None + with trace_tool_span( + tool_name=self._name, + tool_description=self._description, + tool_call_id=call_id, + ): + # Execute the tool's run_stream method + async for result in self.run_stream(self._args_type.model_validate(args), cancellation_token): + return_value = result + yield result + + assert return_value is not None, "The tool must yield a final return value at the end of the stream." + if not isinstance(return_value, self._return_type): + raise TypeError( + f"Expected return value of type {self._return_type.__name__}, but got {type(return_value).__name__}" + ) + + # Log the tool call event + event = ToolCallEvent( + tool_name=self.name, + arguments=dict(args), # Using the raw args passed to run_json + result=self.return_value_as_string(return_value), + ) + logger.info(event) + + class BaseToolWithState(BaseTool[ArgsT, ReturnT], ABC, Generic[ArgsT, ReturnT, StateT], ComponentBase[BaseModel]): def __init__( self, @@ -196,8 +287,8 @@ def save_state(self) -> StateT: ... @abstractmethod def load_state(self, state: StateT) -> None: ... - def save_state_json(self) -> Mapping[str, Any]: + async def save_state_json(self) -> Mapping[str, Any]: return self.save_state().model_dump() - def load_state_json(self, state: Mapping[str, Any]) -> None: + async def load_state_json(self, state: Mapping[str, Any]) -> None: self.load_state(self._state_type.model_validate(state)) diff --git a/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py b/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py index 048b26525901..985d7d1d1201 100644 --- a/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py +++ b/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py @@ -178,4 +178,4 @@ def _from_config(cls, config: FunctionToolConfig) -> Self: if not callable(func): raise TypeError(f"Expected function but got {type(func)}") - return cls(func, "", None) + return cls(func, name=config.name, description=config.description, global_imports=config.global_imports) diff --git a/python/packages/autogen-core/src/autogen_core/tools/_static_workbench.py b/python/packages/autogen-core/src/autogen_core/tools/_static_workbench.py new file mode 100644 index 000000000000..40b1ce47d991 --- /dev/null +++ b/python/packages/autogen-core/src/autogen_core/tools/_static_workbench.py @@ -0,0 +1,225 @@ +import asyncio +import builtins +from typing import Any, AsyncGenerator, Dict, List, Literal, Mapping, Optional + +from pydantic import BaseModel, Field +from typing_extensions import Self + +from .._cancellation_token import CancellationToken +from .._component_config import Component, ComponentModel +from ._base import BaseTool, StreamTool, ToolOverride, ToolSchema +from ._workbench import StreamWorkbench, TextResultContent, ToolResult, Workbench + + +class StaticWorkbenchConfig(BaseModel): + tools: List[ComponentModel] = [] + tool_overrides: Dict[str, ToolOverride] = Field(default_factory=dict) + + +class StateicWorkbenchState(BaseModel): + type: Literal["StaticWorkbenchState"] = "StaticWorkbenchState" + tools: Dict[str, Mapping[str, Any]] = {} + + +class StaticWorkbench(Workbench, Component[StaticWorkbenchConfig]): + """ + A workbench that provides a static set of tools that do not change after + each tool execution. + + Args: + tools (List[BaseTool[Any, Any]]): A list of tools to be included in the workbench. + The tools should be subclasses of :class:`~autogen_core.tools.BaseTool`. + tool_overrides (Optional[Dict[str, ToolOverride]]): Optional mapping of original tool + names to override configurations for name and/or description. This allows + customizing how tools appear to consumers while maintaining the underlying + tool functionality. + """ + + component_provider_override = "autogen_core.tools.StaticWorkbench" + component_config_schema = StaticWorkbenchConfig + + def __init__( + self, tools: List[BaseTool[Any, Any]], tool_overrides: Optional[Dict[str, ToolOverride]] = None + ) -> None: + self._tools = tools + self._tool_overrides = tool_overrides or {} + + # Build reverse mapping from override names to original names for call_tool + self._override_name_to_original: Dict[str, str] = {} + existing_tool_names = {tool.name for tool in self._tools} + + for original_name, override in self._tool_overrides.items(): + if override.name and override.name != original_name: + # Check for conflicts with existing tool names + if override.name in existing_tool_names and override.name != original_name: + raise ValueError( + f"Tool override name '{override.name}' conflicts with existing tool name. " + f"Override names must not conflict with any tool names." + ) + # Check for conflicts with other override names + if override.name in self._override_name_to_original: + existing_original = self._override_name_to_original[override.name] + raise ValueError( + f"Tool override name '{override.name}' is used by multiple tools: " + f"'{existing_original}' and '{original_name}'. Override names must be unique." + ) + self._override_name_to_original[override.name] = original_name + + async def list_tools(self) -> List[ToolSchema]: + result_schemas: List[ToolSchema] = [] + for tool in self._tools: + original_schema = tool.schema + + # Apply overrides if they exist for this tool + if tool.name in self._tool_overrides: + override = self._tool_overrides[tool.name] + # Create a new ToolSchema with overrides applied + schema: ToolSchema = { + "name": override.name if override.name is not None else original_schema["name"], + "description": override.description + if override.description is not None + else original_schema.get("description", ""), + } + # Copy optional fields + if "parameters" in original_schema: + schema["parameters"] = original_schema["parameters"] + if "strict" in original_schema: + schema["strict"] = original_schema["strict"] + else: + schema = original_schema + + result_schemas.append(schema) + return result_schemas + + async def call_tool( + self, + name: str, + arguments: Mapping[str, Any] | None = None, + cancellation_token: CancellationToken | None = None, + call_id: str | None = None, + ) -> ToolResult: + # Check if the name is an override name and map it back to the original + original_name = self._override_name_to_original.get(name, name) + + tool = next((tool for tool in self._tools if tool.name == original_name), None) + if tool is None: + return ToolResult( + name=name, # Return the requested name (which might be overridden) + result=[TextResultContent(content=f"Tool {name} not found.")], + is_error=True, + ) + if not cancellation_token: + cancellation_token = CancellationToken() + if not arguments: + arguments = {} + try: + result_future = asyncio.ensure_future(tool.run_json(arguments, cancellation_token, call_id=call_id)) + cancellation_token.link_future(result_future) + actual_tool_output = await result_future + is_error = False + result_str = tool.return_value_as_string(actual_tool_output) + except Exception as e: + result_str = self._format_errors(e) + is_error = True + return ToolResult(name=name, result=[TextResultContent(content=result_str)], is_error=is_error) + + async def start(self) -> None: + return None + + async def stop(self) -> None: + return None + + async def reset(self) -> None: + return None + + async def save_state(self) -> Mapping[str, Any]: + tool_states = StateicWorkbenchState() + for tool in self._tools: + tool_states.tools[tool.name] = await tool.save_state_json() + return tool_states.model_dump() + + async def load_state(self, state: Mapping[str, Any]) -> None: + parsed_state = StateicWorkbenchState.model_validate(state) + for tool in self._tools: + if tool.name in parsed_state.tools: + await tool.load_state_json(parsed_state.tools[tool.name]) + + def _to_config(self) -> StaticWorkbenchConfig: + return StaticWorkbenchConfig( + tools=[tool.dump_component() for tool in self._tools], tool_overrides=self._tool_overrides + ) + + @classmethod + def _from_config(cls, config: StaticWorkbenchConfig) -> Self: + return cls(tools=[BaseTool.load_component(tool) for tool in config.tools], tool_overrides=config.tool_overrides) + + def _format_errors(self, error: Exception) -> str: + """Recursively format errors into a string.""" + + error_message = "" + if hasattr(builtins, "ExceptionGroup") and isinstance(error, builtins.ExceptionGroup): + # ExceptionGroup is available in Python 3.11+. + # TODO: how to make this compatible with Python 3.10? + for sub_exception in error.exceptions: # type: ignore + error_message += self._format_errors(sub_exception) # type: ignore + else: + error_message += f"{str(error)}\n" + return error_message.strip() + + +class StaticStreamWorkbench(StaticWorkbench, StreamWorkbench): + """ + A workbench that provides a static set of tools that do not change after + each tool execution, and supports streaming results. + """ + + component_provider_override = "autogen_core.tools.StaticStreamWorkbench" + + async def call_tool_stream( + self, + name: str, + arguments: Mapping[str, Any] | None = None, + cancellation_token: CancellationToken | None = None, + call_id: str | None = None, + ) -> AsyncGenerator[Any | ToolResult, None]: + tool = next((tool for tool in self._tools if tool.name == name), None) + if tool is None: + yield ToolResult( + name=name, + result=[TextResultContent(content=f"Tool {name} not found.")], + is_error=True, + ) + return + if not cancellation_token: + cancellation_token = CancellationToken() + if not arguments: + arguments = {} + try: + actual_tool_output: Any | None = None + if isinstance(tool, StreamTool): + previous_result: Any | None = None + try: + async for result in tool.run_json_stream(arguments, cancellation_token, call_id=call_id): + if previous_result is not None: + yield previous_result + previous_result = result + actual_tool_output = previous_result + except Exception as e: + # If there was a previous result before the exception, yield it first + if previous_result is not None: + yield previous_result + # Then yield the error result + result_str = self._format_errors(e) + yield ToolResult(name=tool.name, result=[TextResultContent(content=result_str)], is_error=True) + return + else: + # If the tool is not a stream tool, we run it normally and yield the result + result_future = asyncio.ensure_future(tool.run_json(arguments, cancellation_token, call_id=call_id)) + cancellation_token.link_future(result_future) + actual_tool_output = await result_future + is_error = False + result_str = tool.return_value_as_string(actual_tool_output) + except Exception as e: + result_str = self._format_errors(e) + is_error = True + yield ToolResult(name=tool.name, result=[TextResultContent(content=result_str)], is_error=is_error) diff --git a/python/packages/autogen-core/src/autogen_core/tools/_workbench.py b/python/packages/autogen-core/src/autogen_core/tools/_workbench.py new file mode 100644 index 000000000000..7869c5d4270d --- /dev/null +++ b/python/packages/autogen-core/src/autogen_core/tools/_workbench.py @@ -0,0 +1,216 @@ +from abc import ABC, abstractmethod +from types import TracebackType +from typing import Any, AsyncGenerator, List, Literal, Mapping, Optional, Type + +from pydantic import BaseModel, Field +from typing_extensions import Annotated, Self + +from .._cancellation_token import CancellationToken +from .._component_config import ComponentBase +from .._image import Image +from ._base import ToolSchema + + +class TextResultContent(BaseModel): + """ + Text result content of a tool execution. + """ + + type: Literal["TextResultContent"] = "TextResultContent" + + content: str + """The text content of the result.""" + + +class ImageResultContent(BaseModel): + """ + Image result content of a tool execution. + """ + + type: Literal["ImageResultContent"] = "ImageResultContent" + + content: Image + """The image content of the result.""" + + +ResultContent = Annotated[TextResultContent | ImageResultContent, Field(discriminator="type")] + + +class ToolResult(BaseModel): + """ + A result of a tool execution by a workbench. + """ + + type: Literal["ToolResult"] = "ToolResult" + + name: str + """The name of the tool that was executed.""" + + result: List[ResultContent] + """The result of the tool execution.""" + + is_error: bool = False + """Whether the tool execution resulted in an error.""" + + def to_text(self, replace_image: str | None = None) -> str: + """ + Convert the result to a text string. + + Args: + replace_image (str | None): The string to replace the image content with. + If None, the image content will be included in the text as base64 string. + + Returns: + str: The text representation of the result. + """ + parts: List[str] = [] + for content in self.result: + if isinstance(content, TextResultContent): + parts.append(content.content) + elif isinstance(content, ImageResultContent): + if replace_image is not None: + parts.append(replace_image) + else: + parts.append(f"[Image: {content.content.to_base64()}]") + return "\n".join(parts) + + +class Workbench(ABC, ComponentBase[BaseModel]): + """ + A workbench is a component that provides a set of tools that may share + resources and state. + + A workbench is responsible for managing the lifecycle of the tools and + providing a single interface to call them. The tools provided by the workbench + may be dynamic and their availabilities may change after each tool execution. + + A workbench can be started by calling the :meth:`~autogen_core.tools.Workbench.start` method + and stopped by calling the :meth:`~autogen_core.tools.Workbench.stop` method. + It can also be used as an asynchronous context manager, which will automatically + start and stop the workbench when entering and exiting the context. + """ + + component_type = "workbench" + + @abstractmethod + async def list_tools(self) -> List[ToolSchema]: + """ + List the currently available tools in the workbench as :class:`ToolSchema` + objects. + + The list of tools may be dynamic, and their content may change after + tool execution. + """ + ... + + @abstractmethod + async def call_tool( + self, + name: str, + arguments: Mapping[str, Any] | None = None, + cancellation_token: CancellationToken | None = None, + call_id: str | None = None, + ) -> ToolResult: + """ + Call a tool in the workbench. + + Args: + name (str): The name of the tool to call. + arguments (Mapping[str, Any] | None): The arguments to pass to the tool. + If None, the tool will be called with no arguments. + cancellation_token (CancellationToken | None): An optional cancellation token + to cancel the tool execution. + call_id (str | None): An optional identifier for the tool call, used for tracing. + Returns: + ToolResult: The result of the tool execution. + """ + ... + + @abstractmethod + async def start(self) -> None: + """ + Start the workbench and initialize any resources. + + This method should be called before using the workbench. + """ + ... + + @abstractmethod + async def stop(self) -> None: + """ + Stop the workbench and release any resources. + + This method should be called when the workbench is no longer needed. + """ + ... + + @abstractmethod + async def reset(self) -> None: + """ + Reset the workbench to its initialized, started state. + """ + ... + + @abstractmethod + async def save_state(self) -> Mapping[str, Any]: + """ + Save the state of the workbench. + + This method should be called to persist the state of the workbench. + """ + ... + + @abstractmethod + async def load_state(self, state: Mapping[str, Any]) -> None: + """ + Load the state of the workbench. + + Args: + state (Mapping[str, Any]): The state to load into the workbench. + """ + ... + + async def __aenter__(self) -> Self: + """ + Enter the workbench context manager. + + This method is called when the workbench is used in a `with` statement. + It calls the :meth:`~autogen_core.tools.WorkBench.start` method to start the workbench. + """ + await self.start() + return self + + async def __aexit__( + self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType] + ) -> None: + """ + Exit the workbench context manager. + This method is called when the workbench is used in a `with` statement. + It calls the :meth:`~autogen_core.tools.WorkBench.stop` method to stop the workbench. + """ + await self.stop() + + +class StreamWorkbench(Workbench, ABC): + """A workbench that supports streaming results from tool calls.""" + + @abstractmethod + def call_tool_stream( + self, + name: str, + arguments: Mapping[str, Any] | None = None, + cancellation_token: CancellationToken | None = None, + call_id: str | None = None, + ) -> AsyncGenerator[Any | ToolResult, None]: + """ + Call a tool in the workbench and return a stream of results. + + Args: + name (str): The name of the tool to call. + arguments (Mapping[str, Any] | None): The arguments to pass to the tool + If None, the tool will be called with no arguments. + cancellation_token (CancellationToken | None): An optional cancellation token + to cancel the tool execution. + call_id (str | None): An optional identifier for the tool call, used for tracing. + """ + ... diff --git a/python/packages/autogen-core/src/autogen_core/utils/__init__.py b/python/packages/autogen-core/src/autogen_core/utils/__init__.py new file mode 100644 index 000000000000..e46b75697b81 --- /dev/null +++ b/python/packages/autogen-core/src/autogen_core/utils/__init__.py @@ -0,0 +1,4 @@ +from ._json_to_pydantic import schema_to_pydantic_model +from ._load_json import extract_json_from_str + +__all__ = ["schema_to_pydantic_model", "extract_json_from_str"] diff --git a/python/packages/autogen-core/src/autogen_core/utils/_json_to_pydantic.py b/python/packages/autogen-core/src/autogen_core/utils/_json_to_pydantic.py new file mode 100644 index 000000000000..ba743779e148 --- /dev/null +++ b/python/packages/autogen-core/src/autogen_core/utils/_json_to_pydantic.py @@ -0,0 +1,543 @@ +import datetime +from ipaddress import IPv4Address, IPv6Address +from typing import Annotated, Any, Dict, ForwardRef, List, Literal, Optional, Type, Union, cast + +from pydantic import ( + UUID1, + UUID3, + UUID4, + UUID5, + AnyUrl, + BaseModel, + EmailStr, + Field, + Json, + conbytes, + confloat, + conint, + conlist, + constr, + create_model, +) +from pydantic.fields import FieldInfo + + +class SchemaConversionError(Exception): + """Base class for schema conversion exceptions.""" + + pass + + +class ReferenceNotFoundError(SchemaConversionError): + """Raised when a $ref cannot be resolved.""" + + pass + + +class FormatNotSupportedError(SchemaConversionError): + """Raised when a format is not supported.""" + + pass + + +class UnsupportedKeywordError(SchemaConversionError): + """Raised when an unsupported JSON Schema keyword is encountered.""" + + pass + + +TYPE_MAPPING: Dict[str, Type[Any]] = { + "string": str, + "integer": int, + "boolean": bool, + "number": float, + "array": List, + "object": dict, + "null": type(None), +} + +FORMAT_MAPPING: Dict[str, Any] = { + "uuid": UUID4, + "uuid1": UUID1, + "uuid2": UUID4, + "uuid3": UUID3, + "uuid4": UUID4, + "uuid5": UUID5, + "email": EmailStr, + "uri": AnyUrl, + "hostname": constr(strict=True), + "ipv4": IPv4Address, + "ipv6": IPv6Address, + "ipv4-network": IPv4Address, + "ipv6-network": IPv6Address, + "date-time": datetime.datetime, + "date": datetime.date, + "time": datetime.time, + "duration": datetime.timedelta, + "int32": conint(strict=True, ge=-(2**31), le=2**31 - 1), + "int64": conint(strict=True, ge=-(2**63), le=2**63 - 1), + "float": confloat(strict=True), + "double": float, + "decimal": float, + "byte": conbytes(strict=True), + "binary": conbytes(strict=True), + "password": str, + "path": str, + "json": Json, +} + + +def _make_field( + default: Any, + *, + title: Optional[str] = None, + description: Optional[str] = None, +) -> Any: + """Construct a Pydantic Field with proper typing.""" + field_kwargs: Dict[str, Any] = {} + if title is not None: + field_kwargs["title"] = title + if description is not None: + field_kwargs["description"] = description + return Field(default, **field_kwargs) + + +class _JSONSchemaToPydantic: + def __init__(self) -> None: + self._model_cache: Dict[str, Optional[Union[Type[BaseModel], ForwardRef]]] = {} + + def _resolve_ref(self, ref: str, schema: Dict[str, Any]) -> Dict[str, Any]: + ref_key = ref.split("/")[-1] + definitions = cast(dict[str, dict[str, Any]], schema.get("$defs", {})) + + if ref_key not in definitions: + raise ReferenceNotFoundError( + f"Reference `{ref}` not found in `$defs`. Available keys: {list(definitions.keys())}" + ) + + return definitions[ref_key] + + def get_ref(self, ref_name: str) -> Any: + if ref_name not in self._model_cache: + raise ReferenceNotFoundError( + f"Reference `{ref_name}` not found in cache. Available: {list(self._model_cache.keys())}" + ) + + if self._model_cache[ref_name] is None: + return ForwardRef(ref_name) + + return self._model_cache[ref_name] + + def _process_definitions(self, root_schema: Dict[str, Any]) -> None: + if "$defs" in root_schema: + for model_name in root_schema["$defs"]: + if model_name not in self._model_cache: + self._model_cache[model_name] = None + + for model_name, model_schema in root_schema["$defs"].items(): + if self._model_cache[model_name] is None: + self._model_cache[model_name] = self.json_schema_to_pydantic(model_schema, model_name, root_schema) + + def json_schema_to_pydantic( + self, schema: Dict[str, Any], model_name: str = "GeneratedModel", root_schema: Optional[Dict[str, Any]] = None + ) -> Type[BaseModel]: + if root_schema is None: + root_schema = schema + self._process_definitions(root_schema) + + if "$ref" in schema: + resolved = self._resolve_ref(schema["$ref"], root_schema) + schema = {**resolved, **{k: v for k, v in schema.items() if k != "$ref"}} + + if "allOf" in schema: + merged: Dict[str, Any] = {"type": "object", "properties": {}, "required": []} + for s in schema["allOf"]: + part = self._resolve_ref(s["$ref"], root_schema) if "$ref" in s else s + merged["properties"].update(part.get("properties", {})) + merged["required"].extend(part.get("required", [])) + for k, v in schema.items(): + if k not in {"allOf", "properties", "required"}: + merged[k] = v + merged["required"] = list(set(merged["required"])) + schema = merged + + return self._json_schema_to_model(schema, model_name, root_schema) + + def _resolve_union_types(self, schemas: List[Dict[str, Any]]) -> List[Any]: + types: List[Any] = [] + for s in schemas: + if "$ref" in s: + types.append(self.get_ref(s["$ref"].split("/")[-1])) + elif "enum" in s: + types.append(Literal[tuple(s["enum"])] if len(s["enum"]) > 0 else Any) + else: + json_type = s.get("type") + if json_type not in TYPE_MAPPING: + raise UnsupportedKeywordError(f"Unsupported or missing type `{json_type}` in union") + types.append(TYPE_MAPPING[json_type]) + return types + + def _extract_field_type(self, key: str, value: Dict[str, Any], model_name: str, root_schema: Dict[str, Any]) -> Any: + json_type = value.get("type") + if json_type not in TYPE_MAPPING: + raise UnsupportedKeywordError( + f"Unsupported or missing type `{json_type}` for field `{key}` in `{model_name}`" + ) + + base_type = TYPE_MAPPING[json_type] + constraints: Dict[str, Any] = {} + + if json_type == "string": + if "minLength" in value: + constraints["min_length"] = value["minLength"] + if "maxLength" in value: + constraints["max_length"] = value["maxLength"] + if "pattern" in value: + constraints["pattern"] = value["pattern"] + if constraints: + base_type = constr(**constraints) + + elif json_type == "integer": + if "minimum" in value: + constraints["ge"] = value["minimum"] + if "maximum" in value: + constraints["le"] = value["maximum"] + if "exclusiveMinimum" in value: + constraints["gt"] = value["exclusiveMinimum"] + if "exclusiveMaximum" in value: + constraints["lt"] = value["exclusiveMaximum"] + if constraints: + base_type = conint(**constraints) + + elif json_type == "number": + if "minimum" in value: + constraints["ge"] = value["minimum"] + if "maximum" in value: + constraints["le"] = value["maximum"] + if "exclusiveMinimum" in value: + constraints["gt"] = value["exclusiveMinimum"] + if "exclusiveMaximum" in value: + constraints["lt"] = value["exclusiveMaximum"] + if constraints: + base_type = confloat(**constraints) + + elif json_type == "array": + if "minItems" in value: + constraints["min_length"] = value["minItems"] + if "maxItems" in value: + constraints["max_length"] = value["maxItems"] + item_schema = value.get("items", {"type": "string"}) + if "$ref" in item_schema: + item_type = self.get_ref(item_schema["$ref"].split("/")[-1]) + else: + item_type_name = item_schema.get("type") + if item_type_name is None: + item_type = List[str] + elif item_type_name not in TYPE_MAPPING: + raise UnsupportedKeywordError( + f"Unsupported or missing item type `{item_type_name}` for array field `{key}` in `{model_name}`" + ) + else: + item_type = TYPE_MAPPING[item_type_name] + + base_type = conlist(item_type, **constraints) if constraints else List[item_type] # type: ignore[valid-type] + + if "format" in value: + format_type = FORMAT_MAPPING.get(value["format"]) + if format_type is None: + raise FormatNotSupportedError(f"Unknown format `{value['format']}` for `{key}` in `{model_name}`") + if not isinstance(format_type, type): + return format_type + if not issubclass(format_type, str): + return format_type + return format_type + + return base_type + + def _json_schema_to_model( + self, schema: Dict[str, Any], model_name: str, root_schema: Dict[str, Any] + ) -> Type[BaseModel]: + if "allOf" in schema: + merged: Dict[str, Any] = {"type": "object", "properties": {}, "required": []} + for s in schema["allOf"]: + part = self._resolve_ref(s["$ref"], root_schema) if "$ref" in s else s + merged["properties"].update(part.get("properties", {})) + merged["required"].extend(part.get("required", [])) + for k, v in schema.items(): + if k not in {"allOf", "properties", "required"}: + merged[k] = v + merged["required"] = list(set(merged["required"])) + schema = merged + + fields: Dict[str, tuple[Any, FieldInfo]] = {} + required_fields = set(schema.get("required", [])) + + for key, value in schema.get("properties", {}).items(): + if "$ref" in value: + ref_name = value["$ref"].split("/")[-1] + field_type = self.get_ref(ref_name) + elif "anyOf" in value: + sub_models = self._resolve_union_types(value["anyOf"]) + field_type = Union[tuple(sub_models)] + elif "oneOf" in value: + sub_models = self._resolve_union_types(value["oneOf"]) + field_type = Union[tuple(sub_models)] + if "discriminator" in value: + discriminator = value["discriminator"]["propertyName"] + field_type = Annotated[field_type, Field(discriminator=discriminator)] + elif "enum" in value: + field_type = Literal[tuple(value["enum"])] + elif "allOf" in value: + merged = {"type": "object", "properties": {}, "required": []} + for s in value["allOf"]: + part = self._resolve_ref(s["$ref"], root_schema) if "$ref" in s else s + merged["properties"].update(part.get("properties", {})) + merged["required"].extend(part.get("required", [])) + for k, v in value.items(): + if k not in {"allOf", "properties", "required"}: + merged[k] = v + merged["required"] = list(set(merged["required"])) + field_type = self._json_schema_to_model(merged, f"{model_name}_{key}", root_schema) + elif value.get("type") == "object" and "properties" in value: + field_type = self._json_schema_to_model(value, f"{model_name}_{key}", root_schema) + else: + field_type = self._extract_field_type(key, value, model_name, root_schema) + + if field_type is None: + raise UnsupportedKeywordError(f"Unsupported or missing type for field `{key}` in `{model_name}`") + + default_value = value.get("default") + is_required = key in required_fields + + if not is_required and default_value is None: + field_type = Optional[field_type] + + field_args = { + "default": default_value if not is_required else ..., + } + if "title" in value: + field_args["title"] = value["title"] + if "description" in value: + field_args["description"] = value["description"] + + fields[key] = ( + field_type, + _make_field( + default_value if not is_required else ..., + title=value.get("title"), + description=value.get("description"), + ), + ) + + model: Type[BaseModel] = create_model(model_name, **cast(dict[str, Any], fields)) + model.model_rebuild() + return model + + +def schema_to_pydantic_model(schema: Dict[str, Any], model_name: str = "GeneratedModel") -> Type[BaseModel]: + """ + Convert a JSON Schema dictionary to a fully-typed Pydantic model. + + This function handles schema translation and validation logic to produce + a Pydantic model. + + **Supported JSON Schema Features** + + - **Primitive types**: `string`, `integer`, `number`, `boolean`, `object`, `array`, `null` + - **String formats**: + - `email`, `uri`, `uuid`, `uuid1`, `uuid3`, `uuid4`, `uuid5` + - `hostname`, `ipv4`, `ipv6`, `ipv4-network`, `ipv6-network` + - `date`, `time`, `date-time`, `duration` + - `byte`, `binary`, `password`, `path` + - **String constraints**: + - `minLength`, `maxLength`, `pattern` + - **Numeric constraints**: + - `minimum`, `maximum`, `exclusiveMinimum`, `exclusiveMaximum` + - **Array constraints**: + - `minItems`, `maxItems`, `items` + - **Object schema support**: + - `properties`, `required`, `title`, `description`, `default` + - **Enums**: + - Converted to Python `Literal` type + - **Union types**: + - `anyOf`, `oneOf` supported with optional `discriminator` + - **Inheritance and composition**: + - `allOf` merges multiple schemas into one model + - **$ref and $defs resolution**: + - Supports references to sibling definitions and self-referencing schemas + + .. code-block:: python + + from autogen_core.utils import schema_to_pydantic_model + + # Example 1: Simple user model + schema = { + "title": "User", + "type": "object", + "properties": { + "name": {"type": "string"}, + "email": {"type": "string", "format": "email"}, + "age": {"type": "integer", "minimum": 0}, + }, + "required": ["name", "email"], + } + + UserModel = schema_to_pydantic_model(schema) + user = UserModel(name="Alice", email="alice@example.com", age=30) + + .. code-block:: python + + from autogen_core.utils import schema_to_pydantic_model + + # Example 2: Nested model + schema = { + "title": "BlogPost", + "type": "object", + "properties": { + "title": {"type": "string"}, + "tags": {"type": "array", "items": {"type": "string"}}, + "author": { + "type": "object", + "properties": {"name": {"type": "string"}, "email": {"type": "string", "format": "email"}}, + "required": ["name"], + }, + }, + "required": ["title", "author"], + } + + BlogPost = schema_to_pydantic_model(schema) + + + .. code-block:: python + + from autogen_core.utils import schema_to_pydantic_model + + # Example 3: allOf merging with $refs + schema = { + "title": "EmployeeWithDepartment", + "allOf": [{"$ref": "#/$defs/Employee"}, {"$ref": "#/$defs/Department"}], + "$defs": { + "Employee": { + "type": "object", + "properties": {"id": {"type": "string"}, "name": {"type": "string"}}, + "required": ["id", "name"], + }, + "Department": { + "type": "object", + "properties": {"department": {"type": "string"}}, + "required": ["department"], + }, + }, + } + + Model = schema_to_pydantic_model(schema) + + .. code-block:: python + + from autogen_core.utils import schema_to_pydantic_model + + # Example 4: Self-referencing (recursive) model + schema = { + "title": "Category", + "type": "object", + "properties": { + "name": {"type": "string"}, + "subcategories": {"type": "array", "items": {"$ref": "#/$defs/Category"}}, + }, + "required": ["name"], + "$defs": { + "Category": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "subcategories": {"type": "array", "items": {"$ref": "#/$defs/Category"}}, + }, + "required": ["name"], + } + }, + } + + Category = schema_to_pydantic_model(schema) + + .. code-block:: python + + # Example 5: Serializing and deserializing with Pydantic + + from uuid import uuid4 + from pydantic import BaseModel, EmailStr, Field + from typing import Optional, List, Dict, Any + from autogen_core.utils import schema_to_pydantic_model + + + class Address(BaseModel): + street: str + city: str + zipcode: str + + + class User(BaseModel): + id: str + name: str + email: EmailStr + age: int = Field(..., ge=18) + address: Address + + + class Employee(BaseModel): + id: str + name: str + manager: Optional["Employee"] = None + + + class Department(BaseModel): + name: str + employees: List[Employee] + + + class ComplexModel(BaseModel): + user: User + extra_info: Optional[Dict[str, Any]] = None + sub_items: List[Employee] + + + # Convert ComplexModel to JSON schema + complex_schema = ComplexModel.model_json_schema() + + # Rebuild a new Pydantic model from JSON schema + ReconstructedModel = schema_to_pydantic_model(complex_schema, "ComplexModel") + + # Instantiate reconstructed model + reconstructed = ReconstructedModel( + user={ + "id": str(uuid4()), + "name": "Alice", + "email": "alice@example.com", + "age": 30, + "address": {"street": "123 Main St", "city": "Wonderland", "zipcode": "12345"}, + }, + sub_items=[{"id": str(uuid4()), "name": "Bob", "manager": {"id": str(uuid4()), "name": "Eve"}}], + ) + + print(reconstructed.model_dump()) + + + Args: + schema (Dict[str, Any]): A valid JSON Schema dictionary. + model_name (str, optional): The name of the root model. Defaults to "GeneratedModel". + + Returns: + Type[BaseModel]: A dynamically generated Pydantic model class. + + Raises: + ReferenceNotFoundError: If a `$ref` key references a missing entry. + FormatNotSupportedError: If a `format` keyword is unknown or unsupported. + UnsupportedKeywordError: If the schema contains an unsupported `type`. + + See Also: + - :class:`pydantic.BaseModel` + - :func:`pydantic.create_model` + - https://json-schema.org/ + """ + ... + + return _JSONSchemaToPydantic().json_schema_to_pydantic(schema, model_name) diff --git a/python/packages/autogen-core/src/autogen_core/utils/_load_json.py b/python/packages/autogen-core/src/autogen_core/utils/_load_json.py new file mode 100644 index 000000000000..95ccb0e7c0df --- /dev/null +++ b/python/packages/autogen-core/src/autogen_core/utils/_load_json.py @@ -0,0 +1,20 @@ +import json +import re +from typing import Any, Dict, List + + +def extract_json_from_str(content: str) -> List[Dict[str, Any]]: + """Extract JSON objects from a string. Supports backtick enclosed JSON objects""" + pattern = re.compile(r"```(?:\s*([\w\+\-]+))?\n([\s\S]*?)```") + matches = pattern.findall(content) + ret: List[Dict[str, Any]] = [] + # If no matches found, assume the entire content is a JSON object + if not matches: + ret.append(json.loads(content)) + for match in matches: + language = match[0].strip() if match[0] else None + if language and language.lower() != "json": + raise ValueError(f"Expected JSON object, but found language: {language}") + content = match[1] + ret.append(json.loads(content)) + return ret diff --git a/python/packages/autogen-core/tests/test_base_agent.py b/python/packages/autogen-core/tests/test_base_agent.py index 64bcf59d1774..010bd0624478 100644 --- a/python/packages/autogen-core/tests/test_base_agent.py +++ b/python/packages/autogen-core/tests/test_base_agent.py @@ -9,7 +9,7 @@ async def test_base_agent_create(mocker: MockerFixture) -> None: runtime = mocker.Mock(spec=AgentRuntime) # Shows how to set the context for the agent instantiation in a test context - with AgentInstantiationContext.populate_context((runtime, AgentId("name", "namespace"))): - agent = NoopAgent() - assert agent.runtime == runtime - assert agent.id == AgentId("name", "namespace") + with AgentInstantiationContext.populate_context((runtime, AgentId("name2", "namespace2"))): + agent2 = NoopAgent() + assert agent2.runtime == runtime + assert agent2.id == AgentId("name2", "namespace2") diff --git a/python/packages/autogen-core/tests/test_component_config.py b/python/packages/autogen-core/tests/test_component_config.py index 36125d12828c..9527c19514c4 100644 --- a/python/packages/autogen-core/tests/test_component_config.py +++ b/python/packages/autogen-core/tests/test_component_config.py @@ -361,7 +361,6 @@ async def async_func(x: float, y: float, cancellation_token: CancellationToken) await loaded_async.run_json({"x": 1.0, "y": 2.0}, cancelled_token) -@pytest.mark.asyncio def test_component_descriptions() -> None: """Test different ways of setting component descriptions.""" assert MyComponent("test").dump_component().description is None diff --git a/python/packages/autogen-core/tests/test_json_extraction.py b/python/packages/autogen-core/tests/test_json_extraction.py new file mode 100644 index 000000000000..f511d40977dd --- /dev/null +++ b/python/packages/autogen-core/tests/test_json_extraction.py @@ -0,0 +1,85 @@ +import pytest +from autogen_core.utils import extract_json_from_str + + +def test_extract_json_from_str() -> None: + json_str = """ + { + "name": "John", + "age": 30, + "city": "New York" + } + """ + json_resp = [{"name": "John", "age": 30, "city": "New York"}] + resp = extract_json_from_str(json_str) + assert resp == json_resp + + invalid_json_str = """ + { + "name": "John", + "age": 30, + "city": "New York" + """ + with pytest.raises(ValueError): + extract_json_from_str(invalid_json_str) + + +def test_extract_json_from_str_codeblock() -> None: + code_block_lang_str = """ + ```json + { + "name": "Alice", + "age": 28, + "city": "Seattle" + } + ``` + """ + code_block_no_lang_str = """ + ``` + { + "name": "Alice", + "age": 28, + "city": "Seattle" + } + ``` + """ + code_block_resp = [{"name": "Alice", "age": 28, "city": "Seattle"}] + multi_json_str = """ + ```json + { + "name": "John", + "age": 30, + "city": "New York" + } + ``` + ```json + { + "name": "Jane", + "age": 25, + "city": "Los Angeles" + } + ``` + """ + multi_json_resp = [ + {"name": "John", "age": 30, "city": "New York"}, + {"name": "Jane", "age": 25, "city": "Los Angeles"}, + ] + + lang_resp = extract_json_from_str(code_block_lang_str) + assert lang_resp == code_block_resp + no_lang_resp = extract_json_from_str(code_block_no_lang_str) + assert no_lang_resp == code_block_resp + multi_resp = extract_json_from_str(multi_json_str) + assert multi_resp == multi_json_resp + + invalid_lang_code_block_str = """ + ```notjson + { + "name": "Jane", + "age": 25, + "city": "Los Angeles" + } + ``` + """ + with pytest.raises(ValueError): + extract_json_from_str(invalid_lang_code_block_str) diff --git a/python/packages/autogen-core/tests/test_json_to_pydantic.py b/python/packages/autogen-core/tests/test_json_to_pydantic.py new file mode 100644 index 000000000000..b577ac013759 --- /dev/null +++ b/python/packages/autogen-core/tests/test_json_to_pydantic.py @@ -0,0 +1,760 @@ +import types +from typing import Any, Dict, List, Literal, Optional, Type, get_args, get_origin +from uuid import UUID, uuid4 + +import pytest +from autogen_core.utils._json_to_pydantic import ( + FORMAT_MAPPING, + TYPE_MAPPING, + FormatNotSupportedError, + ReferenceNotFoundError, + UnsupportedKeywordError, + _JSONSchemaToPydantic, # pyright: ignore[reportPrivateUsage] +) +from pydantic import BaseModel, EmailStr, Field, Json, ValidationError + + +# ✅ Define Pydantic models for testing +class Address(BaseModel): + street: str + city: str + zipcode: str + + +class User(BaseModel): + id: UUID + name: str + email: EmailStr + age: int = Field(..., ge=18) # Minimum age = 18 + address: Address + + +class Employee(BaseModel): + id: UUID + name: str + manager: Optional["Employee"] = None # Recursive self-reference + + +class Department(BaseModel): + name: str + employees: List[Employee] # Array of objects + + +class ComplexModel(BaseModel): + user: User + extra_info: Optional[Dict[str, Any]] = None # Optional dictionary + sub_items: List[Employee] # List of Employees + json_string: Optional[Json[Any]] = None # Optional JSON string + + +@pytest.fixture +def converter() -> _JSONSchemaToPydantic: + """Fixture to create a fresh instance of JSONSchemaToPydantic for every test.""" + return _JSONSchemaToPydantic() + + +@pytest.fixture +def sample_json_schema() -> Dict[str, Any]: + """Fixture that returns a JSON schema dynamically using model_json_schema().""" + return User.model_json_schema() + + +@pytest.fixture +def sample_json_schema_recursive() -> Dict[str, Any]: + """Fixture that returns a self-referencing JSON schema.""" + return Employee.model_json_schema() + + +@pytest.fixture +def sample_json_schema_nested() -> Dict[str, Any]: + """Fixture that returns a nested schema with arrays of objects.""" + return Department.model_json_schema() + + +@pytest.fixture +def sample_json_schema_complex() -> Dict[str, Any]: + """Fixture that returns a complex schema with multiple structures.""" + return ComplexModel.model_json_schema() + + +@pytest.mark.parametrize( + "schema_fixture, model_name, expected_fields", + [ + (sample_json_schema, "User", ["id", "name", "email", "age", "address"]), + (sample_json_schema_recursive, "Employee", ["id", "name", "manager"]), + (sample_json_schema_nested, "Department", ["name", "employees"]), + (sample_json_schema_complex, "ComplexModel", ["user", "extra_info", "sub_items", "json_string"]), + ], +) +def test_json_schema_to_pydantic( + converter: _JSONSchemaToPydantic, + schema_fixture: Any, + model_name: str, + expected_fields: List[str], + request: Any, +) -> None: + """Test conversion of JSON Schema to Pydantic model using the class instance.""" + schema = request.getfixturevalue(schema_fixture.__name__) + Model = converter.json_schema_to_pydantic(schema, model_name) + + for field in expected_fields: + assert field in Model.__annotations__, f"Expected '{field}' missing in {model_name}Model" + + +# ✅ **Valid Data Tests** +@pytest.mark.parametrize( + "schema_fixture, model_name, valid_data", + [ + ( + sample_json_schema, + "User", + { + "id": str(uuid4()), + "name": "Alice", + "email": "alice@example.com", + "age": 25, + "address": {"street": "123 Main St", "city": "Metropolis", "zipcode": "12345"}, + }, + ), + ( + sample_json_schema_recursive, + "Employee", + { + "id": str(uuid4()), + "name": "Alice", + "manager": { + "id": str(uuid4()), + "name": "Bob", + }, + }, + ), + ( + sample_json_schema_nested, + "Department", + { + "name": "Engineering", + "employees": [ + { + "id": str(uuid4()), + "name": "Alice", + "manager": { + "id": str(uuid4()), + "name": "Bob", + }, + } + ], + }, + ), + ( + sample_json_schema_complex, + "ComplexModel", + { + "user": { + "id": str(uuid4()), + "name": "Charlie", + "email": "charlie@example.com", + "age": 30, + "address": {"street": "456 Side St", "city": "Gotham", "zipcode": "67890"}, + }, + "extra_info": {"hobby": "Chess", "level": "Advanced"}, + "sub_items": [ + {"id": str(uuid4()), "name": "Eve"}, + {"id": str(uuid4()), "name": "David", "manager": {"id": str(uuid4()), "name": "Frank"}}, + ], + "json_string": '{"foo": "bar"}', + }, + ), + ], +) +def test_valid_data_model( + converter: _JSONSchemaToPydantic, + schema_fixture: Any, + model_name: str, + valid_data: Dict[str, Any], + request: Any, +) -> None: + """Test that valid data is accepted by the generated model.""" + schema = request.getfixturevalue(schema_fixture.__name__) + Model = converter.json_schema_to_pydantic(schema, model_name) + + instance = Model(**valid_data) + assert instance + dumped = instance.model_dump(mode="json", exclude_none=True) + assert dumped == valid_data, f"Model output mismatch.\nExpected: {valid_data}\nGot: {dumped}" + + +# ✅ **Invalid Data Tests** +@pytest.mark.parametrize( + "schema_fixture, model_name, invalid_data", + [ + ( + sample_json_schema, + "User", + { + "id": "not-a-uuid", # Invalid UUID + "name": "Alice", + "email": "not-an-email", # Invalid email + "age": 17, # Below minimum + "address": {"street": "123 Main St", "city": "Metropolis"}, + }, + ), + ( + sample_json_schema_recursive, + "Employee", + { + "id": str(uuid4()), + "name": "Alice", + "manager": { + "id": "not-a-uuid", # Invalid UUID + "name": "Bob", + }, + }, + ), + ( + sample_json_schema_nested, + "Department", + { + "name": "Engineering", + "employees": [ + { + "id": "not-a-uuid", # Invalid UUID + "name": "Alice", + "manager": { + "id": str(uuid4()), + "name": "Bob", + }, + } + ], + }, + ), + ( + sample_json_schema_complex, + "ComplexModel", + { + "user": { + "id": str(uuid4()), + "name": "Charlie", + "email": "charlie@example.com", + "age": "thirty", # Invalid: Should be an int + "address": {"street": "456 Side St", "city": "Gotham", "zipcode": "67890"}, + }, + "extra_info": "should-be-dictionary", # Invalid type + "sub_items": [ + {"id": "invalid-uuid", "name": "Eve"}, # Invalid UUID + {"id": str(uuid4()), "name": 123}, # Invalid name type + ], + "json_string": '{"foo": "bar"', # Invalid JSON + }, + ), + ], +) +def test_invalid_data_model( + converter: _JSONSchemaToPydantic, + schema_fixture: Any, + model_name: str, + invalid_data: Dict[str, Any], + request: Any, +) -> None: + """Test that invalid data raises ValidationError.""" + schema = request.getfixturevalue(schema_fixture.__name__) + Model = converter.json_schema_to_pydantic(schema, model_name) + + with pytest.raises(ValidationError): + Model(**invalid_data) + + +class ListDictModel(BaseModel): + """Example for `List[Dict[str, Any]]`""" + + data: List[Dict[str, Any]] + + +class DictListModel(BaseModel): + """Example for `Dict[str, List[Any]]`""" + + mapping: Dict[str, List[Any]] + + +class NestedListModel(BaseModel): + """Example for `List[List[str]]`""" + + matrix: List[List[str]] + + +@pytest.fixture +def sample_json_schema_list_dict() -> Dict[str, Any]: + """Fixture for `List[Dict[str, Any]]`""" + return ListDictModel.model_json_schema() + + +@pytest.fixture +def sample_json_schema_dict_list() -> Dict[str, Any]: + """Fixture for `Dict[str, List[Any]]`""" + return DictListModel.model_json_schema() + + +@pytest.fixture +def sample_json_schema_nested_list() -> Dict[str, Any]: + """Fixture for `List[List[str]]`""" + return NestedListModel.model_json_schema() + + +@pytest.mark.parametrize( + "schema_fixture, model_name, expected_fields", + [ + (sample_json_schema_list_dict, "ListDictModel", ["data"]), + (sample_json_schema_dict_list, "DictListModel", ["mapping"]), + (sample_json_schema_nested_list, "NestedListModel", ["matrix"]), + ], +) +def test_json_schema_to_pydantic_nested( + converter: _JSONSchemaToPydantic, + schema_fixture: Any, + model_name: str, + expected_fields: list[str], + request: Any, +) -> None: + """Test conversion of JSON Schema to Pydantic model using the class instance.""" + schema = request.getfixturevalue(schema_fixture.__name__) + Model = converter.json_schema_to_pydantic(schema, model_name) + + for field in expected_fields: + assert field in Model.__annotations__, f"Expected '{field}' missing in {model_name}Model" + + +# ✅ **Valid Data Tests** +@pytest.mark.parametrize( + "schema_fixture, model_name, valid_data", + [ + ( + sample_json_schema_list_dict, + "ListDictModel", + { + "data": [ + {"key1": "value1", "key2": 10}, + {"another_key": False, "nested": {"subkey": "data"}}, + ] + }, + ), + ( + sample_json_schema_dict_list, + "DictListModel", + { + "mapping": { + "first": ["a", "b", "c"], + "second": [1, 2, 3, 4], + "third": [True, False, True], + } + }, + ), + ( + sample_json_schema_nested_list, + "NestedListModel", + {"matrix": [["A", "B"], ["C", "D"], ["E", "F"]]}, + ), + ], +) +def test_valid_data_model_nested( + converter: _JSONSchemaToPydantic, + schema_fixture: Any, + model_name: str, + valid_data: Dict[str, Any], + request: Any, +) -> None: + """Test that valid data is accepted by the generated model.""" + schema = request.getfixturevalue(schema_fixture.__name__) + Model = converter.json_schema_to_pydantic(schema, model_name) + + instance = Model(**valid_data) + assert instance + for field, value in valid_data.items(): + assert ( + getattr(instance, field) == value + ), f"Mismatch in field `{field}`: expected `{value}`, got `{getattr(instance, field)}`" + + +# ✅ **Invalid Data Tests** +@pytest.mark.parametrize( + "schema_fixture, model_name, invalid_data", + [ + ( + sample_json_schema_list_dict, + "ListDictModel", + { + "data": "should-be-a-list", # ❌ Should be a list of dicts + }, + ), + ( + sample_json_schema_dict_list, + "DictListModel", + { + "mapping": [ + "should-be-a-dictionary", # ❌ Should be a dict of lists + ] + }, + ), + ( + sample_json_schema_nested_list, + "NestedListModel", + {"matrix": [["A", "B"], "C", ["D", "E"]]}, # ❌ "C" is not a list + ), + ], +) +def test_invalid_data_model_nested( + converter: _JSONSchemaToPydantic, + schema_fixture: Any, + model_name: str, + invalid_data: Dict[str, Any], + request: Any, +) -> None: + """Test that invalid data raises ValidationError.""" + schema = request.getfixturevalue(schema_fixture.__name__) + Model = converter.json_schema_to_pydantic(schema, model_name) + + with pytest.raises(ValidationError): + Model(**invalid_data) + + +def test_reference_not_found(converter: _JSONSchemaToPydantic) -> None: + schema = {"type": "object", "properties": {"manager": {"$ref": "#/$defs/MissingRef"}}} + with pytest.raises(ReferenceNotFoundError): + converter.json_schema_to_pydantic(schema, "MissingRefModel") + + +def test_format_not_supported(converter: _JSONSchemaToPydantic) -> None: + schema = {"type": "object", "properties": {"custom_field": {"type": "string", "format": "unsupported-format"}}} + with pytest.raises(FormatNotSupportedError): + converter.json_schema_to_pydantic(schema, "UnsupportedFormatModel") + + +def test_unsupported_keyword(converter: _JSONSchemaToPydantic) -> None: + schema = {"type": "object", "properties": {"broken_field": {"title": "Missing type"}}} + with pytest.raises(UnsupportedKeywordError): + converter.json_schema_to_pydantic(schema, "MissingTypeModel") + + +def test_enum_field_schema() -> None: + schema = { + "type": "object", + "properties": { + "status": {"type": "string", "enum": ["pending", "approved", "rejected"]}, + "priority": {"type": "integer", "enum": [1, 2, 3]}, + }, + "required": ["status"], + } + + converter: _JSONSchemaToPydantic = _JSONSchemaToPydantic() + Model = converter.json_schema_to_pydantic(schema, "Task") + + status_ann = Model.model_fields["status"].annotation + assert get_origin(status_ann) is Literal + assert set(get_args(status_ann)) == {"pending", "approved", "rejected"} + + priority_ann = Model.model_fields["priority"].annotation + args = get_args(priority_ann) + assert type(None) in args + assert Literal[1, 2, 3] in args + + instance = Model(status="approved", priority=2) + assert instance.status == "approved" # type: ignore[attr-defined] + assert instance.priority == 2 # type: ignore[attr-defined] + + +def test_metadata_title_description(converter: _JSONSchemaToPydantic) -> None: + schema = { + "title": "CustomerProfile", + "description": "A profile containing personal and contact info", + "type": "object", + "properties": { + "first_name": {"type": "string", "title": "First Name", "description": "Given name of the user"}, + "age": {"type": "integer", "title": "Age", "description": "Age in years"}, + "contact": { + "type": "object", + "title": "Contact Information", + "description": "How to reach the user", + "properties": { + "email": { + "type": "string", + "format": "email", + "title": "Email Address", + "description": "Primary email", + } + }, + }, + }, + "required": ["first_name"], + } + + Model: Type[BaseModel] = converter.json_schema_to_pydantic(schema, "CustomerProfile") + generated_schema = Model.model_json_schema() + + assert generated_schema["title"] == "CustomerProfile" + + props = generated_schema["properties"] + assert props["first_name"]["title"] == "First Name" + assert props["first_name"]["description"] == "Given name of the user" + assert props["age"]["title"] == "Age" + assert props["age"]["description"] == "Age in years" + + contact = props["contact"] + assert contact["title"] == "Contact Information" + assert contact["description"] == "How to reach the user" + + # Follow the $ref + ref_key = contact["anyOf"][0]["$ref"].split("/")[-1] + contact_def = generated_schema["$defs"][ref_key] + email = contact_def["properties"]["email"] + assert email["title"] == "Email Address" + assert email["description"] == "Primary email" + + +def test_oneof_with_discriminator(converter: _JSONSchemaToPydantic) -> None: + schema = { + "title": "PetWrapper", + "type": "object", + "properties": { + "pet": { + "oneOf": [{"$ref": "#/$defs/Cat"}, {"$ref": "#/$defs/Dog"}], + "discriminator": {"propertyName": "pet_type"}, + } + }, + "required": ["pet"], + "$defs": { + "Cat": { + "type": "object", + "properties": {"pet_type": {"type": "string", "enum": ["cat"]}, "hunting_skill": {"type": "string"}}, + "required": ["pet_type", "hunting_skill"], + "title": "Cat", + }, + "Dog": { + "type": "object", + "properties": {"pet_type": {"type": "string", "enum": ["dog"]}, "pack_size": {"type": "integer"}}, + "required": ["pet_type", "pack_size"], + "title": "Dog", + }, + }, + } + + Model = converter.json_schema_to_pydantic(schema, "PetWrapper") + + # Instantiate with a Cat + cat = Model(pet={"pet_type": "cat", "hunting_skill": "expert"}) + assert cat.pet.pet_type == "cat" # type: ignore[attr-defined] + + # Instantiate with a Dog + dog = Model(pet={"pet_type": "dog", "pack_size": 4}) + assert dog.pet.pet_type == "dog" # type: ignore[attr-defined] + + # Check round-trip schema includes discriminator + model_schema = Model.model_json_schema() + assert "discriminator" in model_schema["properties"]["pet"] + assert model_schema["properties"]["pet"]["discriminator"]["propertyName"] == "pet_type" + + +def test_allof_merging_with_refs(converter: _JSONSchemaToPydantic) -> None: + schema = { + "title": "EmployeeWithDepartment", + "allOf": [{"$ref": "#/$defs/Employee"}, {"$ref": "#/$defs/Department"}], + "$defs": { + "Employee": { + "type": "object", + "properties": {"id": {"type": "string"}, "name": {"type": "string"}}, + "required": ["id", "name"], + "title": "Employee", + }, + "Department": { + "type": "object", + "properties": {"department": {"type": "string"}}, + "required": ["department"], + "title": "Department", + }, + }, + } + + Model = converter.json_schema_to_pydantic(schema, "EmployeeWithDepartment") + instance = Model(id="123", name="Alice", department="Engineering") + assert instance.id == "123" # type: ignore[attr-defined] + assert instance.name == "Alice" # type: ignore[attr-defined] + assert instance.department == "Engineering" # type: ignore[attr-defined] + + dumped = instance.model_dump() + assert dumped == {"id": "123", "name": "Alice", "department": "Engineering"} + + +def test_nested_allof_merging(converter: _JSONSchemaToPydantic) -> None: + schema = { + "title": "ContainerModel", + "type": "object", + "properties": { + "nested": { + "type": "object", + "properties": { + "data": { + "allOf": [ + {"$ref": "#/$defs/Base"}, + {"type": "object", "properties": {"extra": {"type": "string"}}, "required": ["extra"]}, + ] + } + }, + "required": ["data"], + } + }, + "required": ["nested"], + "$defs": { + "Base": { + "type": "object", + "properties": {"base_field": {"type": "string"}}, + "required": ["base_field"], + "title": "Base", + } + }, + } + + Model = converter.json_schema_to_pydantic(schema, "ContainerModel") + instance = Model(nested={"data": {"base_field": "abc", "extra": "xyz"}}) + + assert instance.nested.data.base_field == "abc" # type: ignore[attr-defined] + assert instance.nested.data.extra == "xyz" # type: ignore[attr-defined] + + +@pytest.mark.parametrize( + "schema, field_name, valid_values, invalid_values", + [ + # String constraints + ( + { + "type": "object", + "properties": { + "username": {"type": "string", "minLength": 3, "maxLength": 10, "pattern": "^[a-zA-Z0-9_]+$"} + }, + "required": ["username"], + }, + "username", + ["user_123", "abc", "Name2023"], + ["", "ab", "toolongusername123", "invalid!char"], + ), + # Integer constraints + ( + { + "type": "object", + "properties": {"age": {"type": "integer", "minimum": 18, "maximum": 99}}, + "required": ["age"], + }, + "age", + [18, 25, 99], + [17, 100, -1], + ), + # Float constraints + ( + { + "type": "object", + "properties": {"score": {"type": "number", "minimum": 0.0, "exclusiveMaximum": 1.0}}, + "required": ["score"], + }, + "score", + [0.0, 0.5, 0.999], + [-0.1, 1.0, 2.5], + ), + # Array constraints + ( + { + "type": "object", + "properties": {"tags": {"type": "array", "items": {"type": "string"}, "minItems": 1, "maxItems": 3}}, + "required": ["tags"], + }, + "tags", + [["a"], ["a", "b"], ["x", "y", "z"]], + [[], ["one", "two", "three", "four"]], + ), + ], +) +def test_field_constraints( + schema: Dict[str, Any], + field_name: str, + valid_values: List[Any], + invalid_values: List[Any], +) -> None: + converter = _JSONSchemaToPydantic() + Model = converter.json_schema_to_pydantic(schema, "ConstraintModel") + + for value in valid_values: + instance = Model(**{field_name: value}) + assert getattr(instance, field_name) == value + + for value in invalid_values: + with pytest.raises(ValidationError): + Model(**{field_name: value}) + + +@pytest.mark.parametrize( + "schema", + [ + # Top-level field + {"type": "object", "properties": {"weird": {"type": "abc"}}, "required": ["weird"]}, + # Inside array items + {"type": "object", "properties": {"items": {"type": "array", "items": {"type": "abc"}}}, "required": ["items"]}, + # Inside anyOf + { + "type": "object", + "properties": {"choice": {"anyOf": [{"type": "string"}, {"type": "abc"}]}}, + "required": ["choice"], + }, + ], +) +def test_unknown_type_raises(schema: Dict[str, Any]) -> None: + converter = _JSONSchemaToPydantic() + with pytest.raises(UnsupportedKeywordError): + converter.json_schema_to_pydantic(schema, "UnknownTypeModel") + + +@pytest.mark.parametrize("json_type, expected_type", list(TYPE_MAPPING.items())) +def test_basic_type_mapping(json_type: str, expected_type: type) -> None: + schema = { + "type": "object", + "properties": {"field": {"type": json_type}}, + "required": ["field"], + } + converter = _JSONSchemaToPydantic() + Model = converter.json_schema_to_pydantic(schema, f"{json_type.capitalize()}Model") + + assert "field" in Model.__annotations__ + field_type = Model.__annotations__["field"] + + # For array/object/null we check the outer type only + if json_type == "null": + assert field_type is type(None) + elif json_type == "array": + assert getattr(field_type, "__origin__", None) is list + elif json_type == "object": + assert field_type in (dict, Dict) or getattr(field_type, "__origin__", None) in (dict, Dict) + + else: + assert field_type == expected_type + + +@pytest.mark.parametrize("format_name, expected_type", list(FORMAT_MAPPING.items())) +def test_format_mapping(format_name: str, expected_type: Any) -> None: + schema = { + "type": "object", + "properties": {"field": {"type": "string", "format": format_name}}, + "required": ["field"], + } + converter = _JSONSchemaToPydantic() + Model = converter.json_schema_to_pydantic(schema, f"{format_name.capitalize()}Model") + + assert "field" in Model.__annotations__ + field_type = Model.__annotations__["field"] + if isinstance(expected_type, types.FunctionType): # if it's a constrained constructor (e.g., conint) + assert callable(field_type) + else: + assert field_type == expected_type + + +def test_unknown_format_raises() -> None: + schema = { + "type": "object", + "properties": {"bad_field": {"type": "string", "format": "definitely-not-a-format"}}, + } + converter = _JSONSchemaToPydantic() + with pytest.raises(FormatNotSupportedError): + converter.json_schema_to_pydantic(schema, "UnknownFormatModel") diff --git a/python/packages/autogen-core/tests/test_model_context.py b/python/packages/autogen-core/tests/test_model_context.py index 46f4b6319370..7a37c30a220b 100644 --- a/python/packages/autogen-core/tests/test_model_context.py +++ b/python/packages/autogen-core/tests/test_model_context.py @@ -4,9 +4,18 @@ from autogen_core.model_context import ( BufferedChatCompletionContext, HeadAndTailChatCompletionContext, + TokenLimitedChatCompletionContext, UnboundedChatCompletionContext, ) -from autogen_core.models import AssistantMessage, LLMMessage, UserMessage +from autogen_core.models import ( + AssistantMessage, + ChatCompletionClient, + FunctionExecutionResultMessage, + LLMMessage, + UserMessage, +) +from autogen_ext.models.ollama import OllamaChatCompletionClient +from autogen_ext.models.openai import OpenAIChatCompletionClient @pytest.mark.asyncio @@ -104,3 +113,95 @@ async def test_unbounded_model_context() -> None: retrieved = await model_context.get_messages() assert len(retrieved) == 3 assert retrieved == messages + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_client,token_limit", + [ + (OpenAIChatCompletionClient(model="gpt-4.1-nano", temperature=0.0, api_key="test"), 30), + (OllamaChatCompletionClient(model="llama3.3"), 20), + ], + ids=["openai", "ollama"], +) +async def test_token_limited_model_context_with_token_limit( + model_client: ChatCompletionClient, token_limit: int +) -> None: + model_context = TokenLimitedChatCompletionContext(model_client=model_client, token_limit=token_limit) + messages: List[LLMMessage] = [ + UserMessage(content="Hello!", source="user"), + AssistantMessage(content="What can I do for you?", source="assistant"), + UserMessage(content="Tell what are some fun things to do in seattle.", source="user"), + ] + for msg in messages: + await model_context.add_message(msg) + + retrieved = await model_context.get_messages() + assert len(retrieved) == 1 # Token limit set very low, will remove 2 of the messages + assert retrieved != messages # Will not be equal to the original messages + + await model_context.clear() + retrieved = await model_context.get_messages() + assert len(retrieved) == 0 + + # Test saving and loading state. + for msg in messages: + await model_context.add_message(msg) + state = await model_context.save_state() + await model_context.clear() + await model_context.load_state(state) + retrieved = await model_context.get_messages() + assert len(retrieved) == 1 + assert retrieved != messages + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_client", + [ + OpenAIChatCompletionClient(model="gpt-4.1-nano", temperature=0.0, api_key="test_key"), + OllamaChatCompletionClient(model="llama3.3"), + ], + ids=["openai", "ollama"], +) +async def test_token_limited_model_context_without_token_limit(model_client: ChatCompletionClient) -> None: + model_context = TokenLimitedChatCompletionContext(model_client=model_client) + messages: List[LLMMessage] = [ + UserMessage(content="Hello!", source="user"), + AssistantMessage(content="What can I do for you?", source="assistant"), + UserMessage(content="Tell what are some fun things to do in seattle.", source="user"), + ] + for msg in messages: + await model_context.add_message(msg) + + retrieved = await model_context.get_messages() + assert len(retrieved) == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_client,token_limit", + [ + (OpenAIChatCompletionClient(model="gpt-4.1-nano", temperature=0.0, api_key="test"), 60), + (OllamaChatCompletionClient(model="llama3.3"), 50), + ], + ids=["openai", "ollama"], +) +async def test_token_limited_model_context_openai_with_function_result( + model_client: ChatCompletionClient, token_limit: int +) -> None: + model_context = TokenLimitedChatCompletionContext(model_client=model_client, token_limit=token_limit) + messages: List[LLMMessage] = [ + FunctionExecutionResultMessage(content=[]), + UserMessage(content="Hello!", source="user"), + AssistantMessage(content="What can I do for you?", source="assistant"), + UserMessage(content="Tell what are some fun things to do in seattle.", source="user"), + ] + for msg in messages: + await model_context.add_message(msg) + + retrieved = await model_context.get_messages() + assert len(retrieved) == 3 # Token limit set very low, will remove 1 of the messages + assert type(retrieved[0]) == UserMessage # Function result should be removed + assert type(retrieved[1]) == AssistantMessage + assert type(retrieved[2]) == UserMessage diff --git a/python/packages/autogen-core/tests/test_runtime.py b/python/packages/autogen-core/tests/test_runtime.py index 64a1cccf4b12..e93a57a6a291 100644 --- a/python/packages/autogen-core/tests/test_runtime.py +++ b/python/packages/autogen-core/tests/test_runtime.py @@ -82,6 +82,60 @@ def agent_factory() -> NoopAgent: await runtime.register_factory(type=AgentType("name2"), agent_factory=agent_factory, expected_class=NoopAgent) +@pytest.mark.asyncio +async def test_agent_type_register_instance() -> None: + runtime = SingleThreadedAgentRuntime() + agent1_id = AgentId(type="name", key="default") + agent2_id = AgentId(type="name", key="notdefault") + agent1 = NoopAgent() + agent1_dup = NoopAgent() + agent2 = NoopAgent() + await agent1.register_instance(runtime=runtime, agent_id=agent1_id) + await agent2.register_instance(runtime=runtime, agent_id=agent2_id) + + assert await runtime.try_get_underlying_agent_instance(agent1_id, type=NoopAgent) == agent1 + assert await runtime.try_get_underlying_agent_instance(agent2_id, type=NoopAgent) == agent2 + with pytest.raises(ValueError): + await agent1_dup.register_instance(runtime=runtime, agent_id=agent1_id) + + +@pytest.mark.asyncio +async def test_agent_type_register_instance_different_types() -> None: + runtime = SingleThreadedAgentRuntime() + agent_id1 = AgentId(type="name", key="noop") + agent_id2 = AgentId(type="name", key="loopback") + agent1 = NoopAgent() + agent2 = LoopbackAgent() + await agent1.register_instance(runtime=runtime, agent_id=agent_id1) + with pytest.raises(ValueError): + await agent2.register_instance(runtime=runtime, agent_id=agent_id2) + + +@pytest.mark.asyncio +async def test_agent_type_register_instance_publish_new_source() -> None: + runtime = SingleThreadedAgentRuntime(ignore_unhandled_exceptions=False) + agent_id = AgentId(type="name", key="default") + agent1 = LoopbackAgent() + await agent1.register_instance(runtime=runtime, agent_id=agent_id) + await runtime.add_subscription(TypeSubscription("notdefault", "name")) + + runtime.start() + with pytest.raises(RuntimeError): + await runtime.publish_message(MessageType(), TopicId("notdefault", "notdefault")) + await runtime.stop_when_idle() + await runtime.close() + + +@pytest.mark.asyncio +async def test_register_instance_factory() -> None: + runtime = SingleThreadedAgentRuntime() + agent1_id = AgentId(type="name", key="default") + agent1 = NoopAgent() + await agent1.register_instance(runtime=runtime, agent_id=agent1_id) + with pytest.raises(ValueError): + await NoopAgent.register(runtime, "name", lambda: NoopAgent()) + + @pytest.mark.asyncio async def test_register_receives_publish(tracer_provider: TracerProvider) -> None: runtime = SingleThreadedAgentRuntime(tracer_provider=tracer_provider) diff --git a/python/packages/autogen-core/tests/test_static_workbench_overrides.py b/python/packages/autogen-core/tests/test_static_workbench_overrides.py new file mode 100644 index 000000000000..37cf1b752f35 --- /dev/null +++ b/python/packages/autogen-core/tests/test_static_workbench_overrides.py @@ -0,0 +1,285 @@ +from typing import Annotated, Dict + +import pytest +from autogen_core.code_executor import ImportFromModule +from autogen_core.tools import FunctionTool, StaticWorkbench, ToolOverride, Workbench + + +@pytest.mark.asyncio +async def test_static_workbench_with_tool_overrides() -> None: + """Test StaticWorkbench with tool name and description overrides.""" + + def test_tool_func_1(x: Annotated[int, "The number to double."]) -> int: + return x * 2 + + def test_tool_func_2(a: Annotated[int, "First number"], b: Annotated[int, "Second number"]) -> int: + return a + b + + test_tool_1 = FunctionTool( + test_tool_func_1, + name="double", + description="A test tool that doubles a number.", + global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])], + ) + test_tool_2 = FunctionTool( + test_tool_func_2, + name="add", + description="A test tool that adds two numbers.", + global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])], + ) + + # Define tool overrides + overrides: Dict[str, ToolOverride] = { + "double": ToolOverride(name="multiply_by_two", description="Multiplies a number by 2"), + "add": ToolOverride(description="Performs addition of two integers"), # Only override description + } + + # Create a StaticWorkbench instance with tool overrides + async with StaticWorkbench(tools=[test_tool_1, test_tool_2], tool_overrides=overrides) as workbench: + # List tools and verify overrides are applied + tools = await workbench.list_tools() + assert len(tools) == 2 + + # Check first tool has name and description overridden + assert tools[0]["name"] == "multiply_by_two" + assert tools[0].get("description") == "Multiplies a number by 2" + assert tools[0].get("parameters") == { + "type": "object", + "properties": {"x": {"type": "integer", "title": "X", "description": "The number to double."}}, + "required": ["x"], + "additionalProperties": False, + } + + # Check second tool has only description overridden + assert tools[1]["name"] == "add" # Original name + assert tools[1].get("description") == "Performs addition of two integers" # Overridden description + assert tools[1].get("parameters") == { + "type": "object", + "properties": { + "a": {"type": "integer", "title": "A", "description": "First number"}, + "b": {"type": "integer", "title": "B", "description": "Second number"}, + }, + "required": ["a", "b"], + "additionalProperties": False, + } + + # Call tools using override names + result_1 = await workbench.call_tool("multiply_by_two", {"x": 5}) + assert result_1.name == "multiply_by_two" # Should return the override name + assert result_1.result[0].type == "TextResultContent" + assert result_1.result[0].content == "10" + assert result_1.to_text() == "10" + assert result_1.is_error is False + + # Call tool using original name (should still work for description-only override) + result_2 = await workbench.call_tool("add", {"a": 3, "b": 7}) + assert result_2.name == "add" + assert result_2.result[0].type == "TextResultContent" + assert result_2.result[0].content == "10" + assert result_2.to_text() == "10" + assert result_2.is_error is False + + # Test calling non-existent tool + result_3 = await workbench.call_tool("nonexistent", {"x": 5}) + assert result_3.name == "nonexistent" + assert result_3.is_error is True + assert result_3.result[0].type == "TextResultContent" + assert "Tool nonexistent not found" in result_3.result[0].content + + +@pytest.mark.asyncio +async def test_static_workbench_without_overrides() -> None: + """Test StaticWorkbench without overrides (original behavior).""" + + def test_tool_func(x: Annotated[int, "The number to double."]) -> int: + return x * 2 + + test_tool = FunctionTool( + test_tool_func, + name="double", + description="A test tool that doubles a number.", + global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])], + ) + + # Create workbench without overrides + async with StaticWorkbench(tools=[test_tool]) as workbench: + tools = await workbench.list_tools() + assert len(tools) == 1 + assert tools[0].get("name") == "double" + assert tools[0].get("description") == "A test tool that doubles a number." + + +@pytest.mark.asyncio +async def test_static_workbench_serialization_with_overrides() -> None: + """Test that StaticWorkbench can be serialized and deserialized with overrides.""" + + def test_tool_func(x: Annotated[int, "The number to double."]) -> int: + return x * 2 + + test_tool = FunctionTool( + test_tool_func, + name="double", + description="A test tool that doubles a number.", + global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])], + ) + + overrides: Dict[str, ToolOverride] = { + "double": ToolOverride(name="multiply_by_two", description="Multiplies a number by 2") + } + + # Create workbench with overrides + workbench = StaticWorkbench(tools=[test_tool], tool_overrides=overrides) + + # Save configuration + config = workbench.dump_component() + assert "tool_overrides" in config.config + + # Load workbench from configuration + async with Workbench.load_component(config) as new_workbench: + tools = await new_workbench.list_tools() + assert len(tools) == 1 + assert tools[0]["name"] == "multiply_by_two" + assert tools[0].get("description") == "Multiplies a number by 2" + + # Test calling tool with override name + result = await new_workbench.call_tool("multiply_by_two", {"x": 5}) + assert result.name == "multiply_by_two" + assert result.result[0].content == "10" + assert result.is_error is False + + +@pytest.mark.asyncio +async def test_static_workbench_partial_overrides() -> None: + """Test StaticWorkbench with partial overrides (name only, description only).""" + + def tool1_func(x: Annotated[int, "Number"]) -> int: + return x + + def tool2_func(x: Annotated[int, "Number"]) -> int: + return x + + tool1 = FunctionTool( + tool1_func, + name="tool1", + description="Original description 1", + global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])], + ) + tool2 = FunctionTool( + tool2_func, + name="tool2", + description="Original description 2", + global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])], + ) + + overrides: Dict[str, ToolOverride] = { + "tool1": ToolOverride(name="renamed_tool1"), # Only name override + "tool2": ToolOverride(description="New description 2"), # Only description override + } + + async with StaticWorkbench(tools=[tool1, tool2], tool_overrides=overrides) as workbench: + tools = await workbench.list_tools() + + # tool1: name overridden, description unchanged + assert tools[0].get("name") == "renamed_tool1" + assert tools[0].get("description") == "Original description 1" + + # tool2: name unchanged, description overridden + assert tools[1].get("name") == "tool2" + assert tools[1].get("description") == "New description 2" + + # Test calling with override name + result1 = await workbench.call_tool("renamed_tool1", {"x": 42}) + assert result1.name == "renamed_tool1" + assert result1.result[0].content == "42" + + # Test calling with original name + result2 = await workbench.call_tool("tool2", {"x": 42}) + assert result2.name == "tool2" + assert result2.result[0].content == "42" + + +def test_tool_override_model() -> None: + """Test ToolOverride model functionality.""" + + # Test with both fields + override1 = ToolOverride(name="new_name", description="new_desc") + assert override1.name == "new_name" + assert override1.description == "new_desc" + + # Test with only name + override2 = ToolOverride(name="new_name") + assert override2.name == "new_name" + assert override2.description is None + + # Test with only description + override3 = ToolOverride(description="new_desc") + assert override3.name is None + assert override3.description == "new_desc" + + # Test empty + override4 = ToolOverride() + assert override4.name is None + assert override4.description is None + + +def test_static_workbench_conflict_detection() -> None: + """Test that StaticWorkbench detects conflicts in tool override names.""" + + def test_tool_func_1(x: Annotated[int, "Number"]) -> int: + return x + + def test_tool_func_2(x: Annotated[int, "Number"]) -> int: + return x + + def test_tool_func_3(x: Annotated[int, "Number"]) -> int: + return x + + tool1 = FunctionTool( + test_tool_func_1, + name="tool1", + description="Tool 1", + global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])], + ) + tool2 = FunctionTool( + test_tool_func_2, + name="tool2", + description="Tool 2", + global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])], + ) + tool3 = FunctionTool( + test_tool_func_3, + name="tool3", + description="Tool 3", + global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])], + ) + + # Test 1: Valid overrides - should work + overrides_valid: Dict[str, ToolOverride] = { + "tool1": ToolOverride(name="renamed_tool1"), + "tool2": ToolOverride(name="renamed_tool2"), + } + workbench_valid = StaticWorkbench(tools=[tool1, tool2, tool3], tool_overrides=overrides_valid) + assert "renamed_tool1" in workbench_valid._override_name_to_original # type: ignore[reportPrivateUsage] + assert "renamed_tool2" in workbench_valid._override_name_to_original # type: ignore[reportPrivateUsage] + + # Test 2: Conflict with existing tool name - should fail + overrides_conflict: Dict[str, ToolOverride] = { + "tool1": ToolOverride(name="tool2") # tool2 already exists + } + with pytest.raises(ValueError): + StaticWorkbench(tools=[tool1, tool2, tool3], tool_overrides=overrides_conflict) + + # Test 3: Duplicate override names - should fail + overrides_duplicate: Dict[str, ToolOverride] = { + "tool1": ToolOverride(name="same_name"), + "tool2": ToolOverride(name="same_name"), # Duplicate + } + with pytest.raises(ValueError): + StaticWorkbench(tools=[tool1, tool2, tool3], tool_overrides=overrides_duplicate) + + # Test 4: Self-renaming - should work but not add to reverse mapping + overrides_self: Dict[str, ToolOverride] = { + "tool1": ToolOverride(name="tool1") # Renaming to itself + } + workbench_self = StaticWorkbench(tools=[tool1, tool2, tool3], tool_overrides=overrides_self) + assert "tool1" not in workbench_self._override_name_to_original # type: ignore[reportPrivateUsage] diff --git a/python/packages/autogen-core/tests/test_tool_agent.py b/python/packages/autogen-core/tests/test_tool_agent.py index 727be91f6707..edb2acfbf375 100644 --- a/python/packages/autogen-core/tests/test_tool_agent.py +++ b/python/packages/autogen-core/tests/test_tool_agent.py @@ -1,7 +1,7 @@ import asyncio import json import logging -from typing import Any, AsyncGenerator, List, Mapping, Optional, Sequence, Union +from typing import Any, AsyncGenerator, List, Literal, Mapping, Optional, Sequence, Union import pytest from autogen_core import EVENT_LOGGER_NAME, AgentId, CancellationToken, FunctionCall, SingleThreadedAgentRuntime @@ -102,6 +102,7 @@ async def create( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -127,6 +128,7 @@ def create_stream( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, diff --git a/python/packages/autogen-core/tests/test_workbench.py b/python/packages/autogen-core/tests/test_workbench.py new file mode 100644 index 000000000000..59f24b1c46b4 --- /dev/null +++ b/python/packages/autogen-core/tests/test_workbench.py @@ -0,0 +1,337 @@ +from typing import Annotated, AsyncGenerator + +import pytest +from autogen_core._cancellation_token import CancellationToken +from autogen_core.code_executor import ImportFromModule +from autogen_core.tools import ( + BaseStreamTool, + FunctionTool, + StaticStreamWorkbench, + StaticWorkbench, + TextResultContent, + ToolResult, + Workbench, +) +from pydantic import BaseModel + + +class StreamArgs(BaseModel): + count: int + + +class StreamResult(BaseModel): + final_count: int + + +class StreamItem(BaseModel): + current: int + + +class StreamTool(BaseStreamTool[StreamArgs, StreamItem, StreamResult]): + def __init__(self) -> None: + super().__init__( + args_type=StreamArgs, + return_type=StreamResult, + name="test_stream_tool", + description="A test stream tool that counts up to a number.", + ) + + async def run(self, args: StreamArgs, cancellation_token: CancellationToken) -> StreamResult: + # For the regular run method, just return the final result + return StreamResult(final_count=args.count) + + async def run_stream( + self, args: StreamArgs, cancellation_token: CancellationToken + ) -> AsyncGenerator[StreamItem | StreamResult, None]: + for i in range(1, args.count + 1): + if cancellation_token.is_cancelled(): + break + yield StreamItem(current=i) + yield StreamResult(final_count=args.count) + + +class StreamToolWithError(BaseStreamTool[StreamArgs, StreamItem, StreamResult]): + def __init__(self) -> None: + super().__init__( + args_type=StreamArgs, + return_type=StreamResult, + name="test_stream_tool_error", + description="A test stream tool that raises an error.", + ) + + async def run(self, args: StreamArgs, cancellation_token: CancellationToken) -> StreamResult: + # For the regular run method, just raise the error + raise ValueError("Stream tool error") + + async def run_stream( + self, args: StreamArgs, cancellation_token: CancellationToken + ) -> AsyncGenerator[StreamItem | StreamResult, None]: + yield StreamItem(current=1) + raise ValueError("Stream tool error") + + +@pytest.mark.asyncio +async def test_static_workbench() -> None: + def test_tool_func_1(x: Annotated[int, "The number to double."]) -> int: + return x * 2 + + def test_tool_func_2(x: Annotated[int, "The number to add 2."]) -> int: + raise ValueError("This is a test error") # Simulate an error + + test_tool_1 = FunctionTool( + test_tool_func_1, + name="test_tool_1", + description="A test tool that doubles a number.", + global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])], + ) + test_tool_2 = FunctionTool( + test_tool_func_2, + name="test_tool_2", + description="A test tool that adds 2 to a number.", + global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])], + ) + + # Create a StaticWorkbench instance with the test tools. + async with StaticWorkbench(tools=[test_tool_1, test_tool_2]) as workbench: + # List tools + tools = await workbench.list_tools() + assert len(tools) == 2 + assert "description" in tools[0] + assert "parameters" in tools[0] + assert tools[0]["name"] == "test_tool_1" + assert tools[0]["description"] == "A test tool that doubles a number." + assert tools[0]["parameters"] == { + "type": "object", + "properties": {"x": {"type": "integer", "title": "X", "description": "The number to double."}}, + "required": ["x"], + "additionalProperties": False, + } + assert "description" in tools[1] + assert "parameters" in tools[1] + assert tools[1]["name"] == "test_tool_2" + assert tools[1]["description"] == "A test tool that adds 2 to a number." + assert tools[1]["parameters"] == { + "type": "object", + "properties": {"x": {"type": "integer", "title": "X", "description": "The number to add 2."}}, + "required": ["x"], + "additionalProperties": False, + } + + # Call tools + result_1 = await workbench.call_tool("test_tool_1", {"x": 5}) + assert result_1.name == "test_tool_1" + assert result_1.result[0].type == "TextResultContent" + assert result_1.result[0].content == "10" + assert result_1.to_text() == "10" + assert result_1.is_error is False + + # Call tool with error + result_2 = await workbench.call_tool("test_tool_2", {"x": 5}) + assert result_2.name == "test_tool_2" + assert result_2.result[0].type == "TextResultContent" + assert result_2.result[0].content == "This is a test error" + assert result_2.to_text() == "This is a test error" + assert result_2.is_error is True + + # Save state. + state = await workbench.save_state() + assert state["type"] == "StaticWorkbenchState" + assert "tools" in state + assert len(state["tools"]) == 2 + + # Dump config. + config = workbench.dump_component() + + # Load the workbench from the config. + async with Workbench.load_component(config) as new_workbench: + # Load state. + await new_workbench.load_state(state) + + # Verify that the tools are still available after loading the state. + tools = await new_workbench.list_tools() + assert len(tools) == 2 + assert "description" in tools[0] + assert "parameters" in tools[0] + assert tools[0]["name"] == "test_tool_1" + assert tools[0]["description"] == "A test tool that doubles a number." + assert tools[0]["parameters"] == { + "type": "object", + "properties": {"x": {"type": "integer", "title": "X", "description": "The number to double."}}, + "required": ["x"], + "additionalProperties": False, + } + assert "description" in tools[1] + assert "parameters" in tools[1] + assert tools[1]["name"] == "test_tool_2" + assert tools[1]["description"] == "A test tool that adds 2 to a number." + assert tools[1]["parameters"] == { + "type": "object", + "properties": {"x": {"type": "integer", "title": "X", "description": "The number to add 2."}}, + "required": ["x"], + "additionalProperties": False, + } + + # Call tools + result_1 = await new_workbench.call_tool("test_tool_1", {"x": 5}) + assert result_1.name == "test_tool_1" + assert result_1.result[0].type == "TextResultContent" + assert result_1.result[0].content == "10" + assert result_1.to_text() == "10" + assert result_1.is_error is False + + # Call tool with error + result_2 = await new_workbench.call_tool("test_tool_2", {"x": 5}) + assert result_2.name == "test_tool_2" + assert result_2.result[0].type == "TextResultContent" + assert result_2.result[0].content == "This is a test error" + assert result_2.to_text() == "This is a test error" + assert result_2.is_error is True + + +@pytest.mark.asyncio +async def test_static_stream_workbench_call_tool_stream() -> None: + """Test call_tool_stream with streaming tools and regular tools.""" + + def regular_tool_func(x: Annotated[int, "The number to double."]) -> int: + return x * 2 + + regular_tool = FunctionTool( + regular_tool_func, + name="regular_tool", + description="A regular tool that doubles a number.", + global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])], + ) + + stream_tool = StreamTool() + stream_tool_with_error = StreamToolWithError() + + async with StaticStreamWorkbench(tools=[regular_tool, stream_tool, stream_tool_with_error]) as workbench: + # Test streaming tool + results: list[StreamItem | StreamResult | ToolResult] = [] + async for result in workbench.call_tool_stream("test_stream_tool", {"count": 3}): + results.append(result) + + # Should get 3 intermediate results and 1 final result + assert len(results) == 4 + + # Check intermediate results (StreamItem objects) + for i, result in enumerate(results[:3]): + assert isinstance(result, StreamItem) + assert result.current == i + 1 + + # Check final result (ToolResult) + final_result = results[-1] + assert isinstance(final_result, ToolResult) + assert final_result.name == "test_stream_tool" + assert final_result.is_error is False + assert final_result.result[0].type == "TextResultContent" + assert "final_count" in final_result.result[0].content + + # Test regular (non-streaming) tool + results_regular: list[ToolResult] = [] + async for result in workbench.call_tool_stream("regular_tool", {"x": 5}): + results_regular.append(result) # type: ignore + + # Should get only 1 result for non-streaming tool + assert len(results_regular) == 1 + final_result = results_regular[0] + assert final_result.name == "regular_tool" + assert final_result.is_error is False + assert final_result.result[0].content == "10" + + # Test streaming tool with error + results_error: list[StreamItem | ToolResult] = [] + async for result in workbench.call_tool_stream("test_stream_tool_error", {"count": 3}): + results_error.append(result) # type: ignore + + # Should get 1 intermediate result and 1 error result + assert len(results_error) == 2 + + # Check intermediate result + intermediate_result = results_error[0] + assert isinstance(intermediate_result, StreamItem) + assert intermediate_result.current == 1 + + # Check error result + error_result = results_error[1] + assert isinstance(error_result, ToolResult) + assert error_result.name == "test_stream_tool_error" + assert error_result.is_error is True + result_content = error_result.result[0] + assert isinstance(result_content, TextResultContent) + assert "Stream tool error" in result_content.content + + # Test tool not found + results_not_found: list[ToolResult] = [] + async for result in workbench.call_tool_stream("nonexistent_tool", {"x": 5}): + results_not_found.append(result) # type: ignore + + assert len(results_not_found) == 1 + error_result = results_not_found[0] + assert error_result.name == "nonexistent_tool" + assert error_result.is_error is True + result_content = error_result.result[0] + assert isinstance(result_content, TextResultContent) + assert "Tool nonexistent_tool not found" in result_content.content + + # Test with no arguments + results_no_args: list[StreamItem | StreamResult | ToolResult] = [] + async for result in workbench.call_tool_stream("test_stream_tool", {"count": 1}): + results_no_args.append(result) # type: ignore + + assert len(results_no_args) == 2 # 1 intermediate + 1 final + + # Test with None arguments + results_none: list[ToolResult] = [] + async for result in workbench.call_tool_stream("regular_tool", None): + results_none.append(result) # type: ignore + + # Should still work but may get error due to missing required argument + assert len(results_none) == 1 + result = results_none[0] + assert result.name == "regular_tool" + # This should error because x is required + assert result.is_error is True + + +@pytest.mark.asyncio +async def test_static_stream_workbench_call_tool_stream_cancellation() -> None: + """Test call_tool_stream with cancellation token.""" + stream_tool = StreamTool() + + async with StaticStreamWorkbench(tools=[stream_tool]) as workbench: + # Test with cancellation token + cancellation_token = CancellationToken() + + results: list[StreamItem | StreamResult | ToolResult] = [] + async for result in workbench.call_tool_stream("test_stream_tool", {"count": 5}, cancellation_token): + results.append(result) # type: ignore + if len(results) == 2: # Cancel after 2 results + cancellation_token.cancel() + + # Should get at least 2 results before cancellation + assert len(results) >= 2 + + +@pytest.mark.asyncio +async def test_static_stream_workbench_inheritance() -> None: + """Test that StaticStreamWorkbench inherits from both StaticWorkbench and StreamWorkbench.""" + stream_tool = StreamTool() + + async with StaticStreamWorkbench(tools=[stream_tool]) as workbench: + # Test that it has regular workbench functionality + tools = await workbench.list_tools() + assert len(tools) == 1 + assert tools[0]["name"] == "test_stream_tool" + + # Test regular call_tool method + result = await workbench.call_tool("test_stream_tool", {"count": 2}) + assert result.name == "test_stream_tool" + assert result.is_error is False + + # Test streaming functionality exists + assert hasattr(workbench, "call_tool_stream") + results: list[StreamItem | StreamResult | ToolResult] = [] + async for result in workbench.call_tool_stream("test_stream_tool", {"count": 2}): + results.append(result) # type: ignore + assert len(results) == 3 # 2 intermediate + 1 final diff --git a/python/packages/autogen-ext/pyproject.toml b/python/packages/autogen-ext/pyproject.toml index 93315d45d784..d68bd0460001 100644 --- a/python/packages/autogen-ext/pyproject.toml +++ b/python/packages/autogen-ext/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "autogen-ext" -version = "0.4.9" +version = "0.7.2" license = {file = "LICENSE-CODE"} description = "AutoGen extensions library" readme = "README.md" @@ -15,22 +15,24 @@ classifiers = [ "Operating System :: OS Independent", ] dependencies = [ - "autogen-core==0.4.9", + "autogen-core==0.7.2", ] [project.optional-dependencies] anthropic = ["anthropic>=0.48"] langchain = ["langchain_core~= 0.3.3"] azure = [ - "azure-ai-inference>=1.0.0b7", + "azure-ai-inference>=1.0.0b9", + "azure-ai-projects>=1.0.0b11", "azure-core", "azure-identity", + "azure-search-documents>=11.4.0", ] docker = ["docker~=7.0", "asyncio_atexit>=1.0.1"] ollama = ["ollama>=0.4.7", "tiktoken>=0.8.0"] -openai = ["openai>=1.66.5", "tiktoken>=0.8.0", "aiofiles"] +openai = ["openai>=1.93", "tiktoken>=0.8.0", "aiofiles"] file-surfer = [ - "autogen-agentchat==0.4.9", + "autogen-agentchat==0.7.2", "magika>=0.6.1rc2", "markitdown[all]~=0.1.0a3", ] @@ -39,24 +41,30 @@ llama-cpp = [ "llama-cpp-python>=0.3.8", ] -graphrag = ["graphrag>=1.0.1"] -chromadb = ["chromadb"] +graphrag = ["graphrag>=2.3.0"] +chromadb = ["chromadb>=1.0.0"] +mem0 = ["mem0ai>=0.1.98"] +mem0-local = [ + "mem0ai>=0.1.98", + "neo4j>=5.25.0", + "chromadb>=1.0.0" +] web-surfer = [ - "autogen-agentchat==0.4.9", + "autogen-agentchat==0.7.2", "playwright>=1.48.0", "pillow>=11.0.0", "magika>=0.6.1rc2", "markitdown[all]~=0.1.0a3", ] magentic-one = [ - "autogen-agentchat==0.4.9", + "autogen-agentchat==0.7.2", "magika>=0.6.1rc2", "markitdown[all]~=0.1.0a3", "playwright>=1.48.0", "pillow>=11.0.0", ] video-surfer = [ - "autogen-agentchat==0.4.9", + "autogen-agentchat==0.7.2", "opencv-python>=4.5", "ffmpeg-python", "openai-whisper", @@ -77,7 +85,15 @@ jupyter-executor = [ "nbclient>=0.10.2", ] -task-centric-memory = ["chromadb>=0.6.3"] +docker-jupyter-executor = [ + "docker~=7.0", + "asyncio_atexit>=1.0.1", + "websockets>=15.0.1", + "requests>=2.32.3", + "aiohttp>=3.11.16", +] + +task-centric-memory = ["chromadb>=1.0.0"] semantic-kernel-core = [ "semantic-kernel>=1.17.1", @@ -134,11 +150,13 @@ semantic-kernel-all = [ rich = ["rich>=13.9.4"] -mcp = [ - "mcp>=1.1.3", - "json-schema-to-pydantic>=0.2.2" +mcp = ["mcp>=1.11.0"] +canvas = [ + "unidiff>=0.7.5", ] +redisvl = ["redisvl>=0.6.0"] + [tool.hatch.build.targets.wheel] packages = ["src/autogen_ext"] @@ -148,6 +166,7 @@ dev = [ "langchain-experimental", "pandas-stubs>=2.2.3.241126", "httpx>=0.28.1", + "opentelemetry-proto>=1.28.0" ] [tool.ruff] @@ -178,7 +197,7 @@ test.sequence = [ test.default_item_type = "cmd" test-grpc = "pytest -n 1 --cov=src --cov-report=term-missing --cov-report=xml --grpc" test-windows = "pytest -n 1 --cov=src --cov-report=term-missing --cov-report=xml -m 'windows'" -mypy = "mypy --config-file ../../pyproject.toml --exclude src/autogen_ext/runtimes/grpc/protos --exclude tests/protos src tests" +mypy = "mypy --config-file ../../pyproject.toml --exclude src/autogen_ext/runtimes/grpc/protos --exclude tests/protos --ignore-missing-imports src tests" [tool.mypy] [[tool.mypy.overrides]] diff --git a/python/packages/autogen-ext/src/autogen_ext/agents/azure/__init__.py b/python/packages/autogen-ext/src/autogen_ext/agents/azure/__init__.py new file mode 100644 index 000000000000..1952155d0413 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/agents/azure/__init__.py @@ -0,0 +1,10 @@ +try: + from ._azure_ai_agent import AzureAIAgent +except ImportError as e: + raise ImportError( + "Dependencies for AzureAIAgent not found. " + 'Please install autogen-ext with the "azure" extra: ' + 'pip install "autogen-ext[azure]"' + ) from e + +__all__ = ["AzureAIAgent"] diff --git a/python/packages/autogen-ext/src/autogen_ext/agents/azure/_azure_ai_agent.py b/python/packages/autogen-ext/src/autogen_ext/agents/azure/_azure_ai_agent.py new file mode 100644 index 000000000000..42b98f81625d --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/agents/azure/_azure_ai_agent.py @@ -0,0 +1,1096 @@ +import asyncio +import json +import logging +import os +from typing import ( + Any, + AsyncGenerator, + Dict, + Iterable, + List, + Mapping, + Optional, + Sequence, + Set, + cast, +) + +from autogen_agentchat import TRACE_LOGGER_NAME +from autogen_agentchat.agents import BaseChatAgent +from autogen_agentchat.base import Response +from autogen_agentchat.messages import ( + AgentEvent, + BaseChatMessage, + ChatMessage, + HandoffMessage, + MultiModalMessage, + StopMessage, + TextMessage, + ToolCallExecutionEvent, + ToolCallRequestEvent, +) +from autogen_core import CancellationToken, FunctionCall +from autogen_core.models._types import FunctionExecutionResult +from autogen_core.tools import FunctionTool, Tool + +from azure.ai.agents.models import ( + Agent, + AgentsResponseFormat, + AgentThread, + AzureAISearchToolDefinition, + AzureFunctionToolDefinition, + BingGroundingToolDefinition, + CodeInterpreterToolDefinition, + CodeInterpreterToolResource, + FileInfo, + FilePurpose, + FileSearchToolDefinition, + FileSearchToolResource, + FileState, + FunctionDefinition, + FunctionToolDefinition, + ListSortOrder, + MessageRole, + MessageTextUrlCitationAnnotation, + RunStatus, + ThreadRun, + ToolDefinition, + ToolOutput, + ToolResources, + VectorStore, + VectorStoreChunkingStrategyRequest, + VectorStoreDataSource, + VectorStoreExpirationPolicy, +) +from azure.ai.agents.models._patch import ThreadMessage +from azure.ai.projects.aio import AIProjectClient + +from ._types import AzureAIAgentState, ListToolType + +trace_logger = logging.getLogger(TRACE_LOGGER_NAME) + + +class AzureAIAgent(BaseChatAgent): + """ + Azure AI Assistant agent for AutoGen. + + Installation: + + .. code-block:: bash + + pip install "autogen-ext[azure]" # For Azure AI Foundry Agent Service + + This agent leverages the Azure AI Assistant API to create AI assistants with capabilities like: + + * Code interpretation and execution + * Grounding with Bing search + * File handling and search + * Custom function calling + * Multi-turn conversations + + The agent integrates with AutoGen's messaging system, providing a seamless way to use Azure AI + capabilities within the AutoGen framework. It supports tools like code interpreter, + file search, and various grounding mechanisms. + + Agent name must be a valid Python identifier: + 1. It must start with a letter (A-Z, a-z) or an underscore (_). + 2. It can only contain letters, digits (0-9), or underscores. + 3. It cannot be a Python keyword. + 4. It cannot contain spaces or special characters. + 5. It cannot start with a digit. + + + Check here on how to create a new secured agent with user-managed identity: + https://learn.microsoft.com/en-us/azure/ai-services/agents/how-to/virtual-networks + + Examples: + + Use the AzureAIAgent to create an agent grounded with Bing: + + .. code-block:: python + + import asyncio + import os + + from autogen_agentchat.messages import TextMessage + from autogen_core import CancellationToken + from autogen_ext.agents.azure._azure_ai_agent import AzureAIAgent + from azure.ai.projects.aio import AIProjectClient + from azure.identity.aio import DefaultAzureCredential + from azure.ai.agents.models import BingGroundingTool + import dotenv + + + async def bing_example(): + async with DefaultAzureCredential() as credential: + async with AIProjectClient( # type: ignore + credential=credential, endpoint=os.getenv("AZURE_PROJECT_ENDPOINT", "") + ) as project_client: + conn = await project_client.connections.get(name=os.getenv("BING_CONNECTION_NAME", "")) + + bing_tool = BingGroundingTool(conn.id) + agent_with_bing_grounding = AzureAIAgent( + name="bing_agent", + description="An AI assistant with Bing grounding", + project_client=project_client, + deployment_name="gpt-4o", + instructions="You are a helpful assistant.", + tools=bing_tool.definitions, + metadata={"source": "AzureAIAgent"}, + ) + + # For the bing grounding tool to return the citations, the message must contain an instruction for the model to do return them. + # For example: "Please provide citations for the answers" + + result = await agent_with_bing_grounding.on_messages( + messages=[ + TextMessage( + content="What is Microsoft\\'s annual leave policy? Provide citations for your answers.", + source="user", + ) + ], + cancellation_token=CancellationToken(), + message_limit=5, + ) + print(result) + + + if __name__ == "__main__": + dotenv.load_dotenv() + asyncio.run(bing_example()) + + Use the AzureAIAgent to create an agent with file search capability: + + .. code-block:: python + + import asyncio + import os + import tempfile + import urllib.request + + import dotenv + from autogen_agentchat.messages import TextMessage + from autogen_core import CancellationToken + from autogen_ext.agents.azure._azure_ai_agent import AzureAIAgent + from azure.ai.projects.aio import AIProjectClient + from azure.identity.aio import DefaultAzureCredential + + + async def file_search_example(): + # Download README.md from GitHub + readme_url = "https://raw.githubusercontent.com/microsoft/autogen/refs/heads/main/README.md" + temp_file = None + + try: + # Create a temporary file to store the downloaded README + temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".md") + urllib.request.urlretrieve(readme_url, temp_file.name) + print(f"Downloaded README.md to {temp_file.name}") + + async with DefaultAzureCredential() as credential: + async with AIProjectClient( # type: ignore + credential=credential, endpoint=os.getenv("AZURE_PROJECT_ENDPOINT", "") + ) as project_client: + agent_with_file_search = AzureAIAgent( + name="file_search_agent", + description="An AI assistant with file search capabilities", + project_client=project_client, + deployment_name="gpt-4.1-mini", + instructions="You are a helpful assistant.", + tools=["file_search"], + metadata={"source": "AzureAIAgent"}, + ) + + ct: CancellationToken = CancellationToken() + # Use the downloaded README file for file search + await agent_with_file_search.on_upload_for_file_search( + file_paths=[temp_file.name], + vector_store_name="file_upload_index", + vector_store_metadata={"source": "AzureAIAgent"}, + cancellation_token=ct, + vector_store_polling_interval=60, + ) + result = await agent_with_file_search.on_messages( + messages=[ + TextMessage( + content="Hello, what is AutoGen and what capabilities does it have?", source="user" + ) + ], + cancellation_token=ct, + message_limit=5, + ) + print(result) + finally: + # Clean up the temporary file + if temp_file and os.path.exists(temp_file.name): + os.unlink(temp_file.name) + print(f"Removed temporary file {temp_file.name}") + + + if __name__ == "__main__": + dotenv.load_dotenv() + asyncio.run(file_search_example()) + + Use the AzureAIAgent to create an agent with code interpreter capability: + + .. code-block:: python + + import asyncio + import os + + import dotenv + from autogen_agentchat.messages import TextMessage + from autogen_core import CancellationToken + from autogen_ext.agents.azure._azure_ai_agent import AzureAIAgent + from azure.ai.projects.aio import AIProjectClient + from azure.identity.aio import DefaultAzureCredential + + + async def code_interpreter_example(): + async with DefaultAzureCredential() as credential: + async with AIProjectClient( # type: ignore + credential=credential, endpoint=os.getenv("AZURE_PROJECT_ENDPOINT", "") + ) as project_client: + agent_with_code_interpreter = AzureAIAgent( + name="code_interpreter_agent", + description="An AI assistant with code interpreter capabilities", + project_client=project_client, + deployment_name="gpt-4.1-mini", + instructions="You are a helpful assistant.", + tools=["code_interpreter"], + metadata={"source": "AzureAIAgent"}, + ) + + await agent_with_code_interpreter.on_upload_for_code_interpreter( + file_paths="/workspaces/autogen/python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/data/nifty_500_quarterly_results.csv", + cancellation_token=CancellationToken(), + polling_interval=5, + ) + + result = await agent_with_code_interpreter.on_messages( + messages=[ + TextMessage( + content="Aggregate the number of stocks per industry and give me a markdown table as a result?", + source="user", + ) + ], + cancellation_token=CancellationToken(), + ) + + print(result) + + + if __name__ == "__main__": + dotenv.load_dotenv() + asyncio.run(code_interpreter_example()) + """ + + def __init__( + self, + name: str, + description: str, + project_client: AIProjectClient, + deployment_name: str, + instructions: str, + tools: Optional[ListToolType] = None, + agent_id: Optional[str] = None, + thread_id: Optional[str] = None, + metadata: Optional[Dict[str, str]] = None, + response_format: Optional[AgentsResponseFormat] = None, + temperature: Optional[float] = None, + tool_resources: Optional[ToolResources] = None, + top_p: Optional[float] = None, + ) -> None: + """ + Initialize the Azure AI Agent. + + Args: + name (str): The name of the agent. Must be a valid Python identifier. + description (str): A brief description of the agent's purpose. + project_client (AIProjectClient): The Azure AI Project client for API interactions. + deployment_name (str): The model deployment name to use for the agent (e.g., "gpt-4"). + instructions (str): Detailed instructions for the agent's behavior. + tools (Optional[Iterable[Union[str, ToolDefinition, Tool, Callable]]]): A list of tools the agent can use. + Supported string values: "file_search", "code_interpreter", "bing_grounding", + "azure_ai_search", "azure_function", "sharepoint_grounding". + agent_id (Optional[str]): Existing agent ID to use instead of creating a new one. + thread_id (Optional[str]): Existing thread ID to continue a conversation. + metadata (Optional[Dict[str, str]]): Additional metadata for the agent. + response_format (Optional[_types.AgentsApiResponseFormatOption]): Format options for the agent's responses. + temperature (Optional[float]): Sampling temperature, controls randomness of output. + tool_resources (Optional[models.ToolResources]): Resources configuration for agent tools. + top_p (Optional[float]): An alternative to temperature, nucleus sampling parameter. + + Raises: + ValueError: If an unsupported tool type is provided. + """ + super().__init__(name, description) + + if tools is None: + tools = [] + + self._original_tools: list[Tool] = [] + + converted_tools: List[ToolDefinition] = [] + self._add_tools(tools, converted_tools) + + self._project_client = project_client + self._agent: Optional[Agent] = None + self._thread: Optional[AgentThread] = None + self._init_thread_id = thread_id + self._deployment_name = deployment_name + self._instructions = instructions + self._api_tools = converted_tools + self._agent_id = agent_id + self._metadata = metadata + self._response_format = response_format + self._temperature = temperature + self._tool_resources = tool_resources + self._top_p = top_p + self._vector_store_id: Optional[str] = None + self._uploaded_file_ids: List[str] = [] + + self._initial_message_ids: Set[str] = set() + self._initial_state_retrieved: bool = False + + # Properties + @property + def produced_message_types(self) -> Sequence[type[ChatMessage]]: + """The types of messages that the assistant agent produces.""" + return (TextMessage,) + + @property + def thread_id(self) -> str: + if self._thread is None: + raise ValueError("Thread not initialized") + return self._thread.id + + @property + def _get_agent_id(self) -> str: + if self._agent is None: + raise ValueError("Agent not initialized") + return self._agent.id + + @property + def description(self) -> str: + if not self._description: + raise ValueError("Description not initialized") + return self._description + + @property + def agent_id(self) -> str: + if not self._agent_id: + raise ValueError("Agent not initialized") + return self._agent_id + + @property + def deployment_name(self) -> str: + if not self._deployment_name: + raise ValueError("Deployment name not initialized") + return self._deployment_name + + @property + def instructions(self) -> str: + if not self._instructions: + raise ValueError("Instructions not initialized") + return self._instructions + + @property + def tools(self) -> List[ToolDefinition]: + """ + Get the list of tools available to the agent. + + Returns: + List[ToolDefinition]: The list of tool definitions. + """ + return self._api_tools + + def _add_tools(self, tools: Optional[ListToolType], converted_tools: List[ToolDefinition]) -> None: + """ + Convert various tool formats to Azure AI Agent tool definitions. + + Args: + tools: List of tools in various formats (string identifiers, ToolDefinition objects, Tool objects, or callables) + converted_tools: List to which converted tool definitions will be added + + Raises: + ValueError: If an unsupported tool type is provided + """ + if tools is None: + return + + for tool in tools: + if isinstance(tool, str): + if tool == "file_search": + converted_tools.append(FileSearchToolDefinition()) + elif tool == "code_interpreter": + converted_tools.append(CodeInterpreterToolDefinition()) + elif tool == "bing_grounding": + converted_tools.append(BingGroundingToolDefinition()) # type: ignore + elif tool == "azure_ai_search": + converted_tools.append(AzureAISearchToolDefinition()) + elif tool == "azure_function": + converted_tools.append(AzureFunctionToolDefinition()) # type: ignore + # elif tool == "sharepoint_grounding": + # converted_tools.append(SharepointToolDefinition()) # type: ignore + else: + raise ValueError(f"Unsupported tool string: {tool}") + elif isinstance(tool, ToolDefinition): + converted_tools.append(tool) + elif isinstance(tool, Tool): + self._original_tools.append(tool) + converted_tools.append(self._convert_tool_to_function_tool_definition(tool)) + elif callable(tool): + if hasattr(tool, "__doc__") and tool.__doc__ is not None: + description = tool.__doc__ + else: + description = "" + function_tool = FunctionTool(tool, description=description) + self._original_tools.append(function_tool) + converted_tools.append(self._convert_tool_to_function_tool_definition(function_tool)) + else: + raise ValueError(f"Unsupported tool type: {type(tool)}") + + def _convert_tool_to_function_tool_definition(self, tool: Tool) -> FunctionToolDefinition: + """ + Convert an autogen Tool to an Azure AI Agent function tool definition. + + Args: + tool (Tool): The AutoGen tool to convert + + Returns: + models.FunctionToolDefinition: A function tool definition compatible with Azure AI Agent API + """ + + schema = tool.schema + parameters: Dict[str, object] = {} + + if "parameters" in schema: + parameters = { + "type": schema["parameters"]["type"], + "properties": schema["parameters"]["properties"], + } + if "required" in schema["parameters"]: + parameters["required"] = schema["parameters"]["required"] + + func_definition = FunctionDefinition(name=tool.name, description=tool.description, parameters=parameters) + + return FunctionToolDefinition( + function=func_definition, + ) + + async def _ensure_initialized(self, create_new_thread: bool = False, create_new_agent: bool = False) -> None: + """ + Ensure agent and thread are properly initialized before operations. + + This method ensures that both the Azure AI Agent and thread are created or retrieved + from existing IDs. It also handles retrieving the initial state of an existing thread + when needed. + + Args: + create_new_thread (bool): When True, creates a new thread even if thread_id is provided + create_new_agent (bool): When True, creates a new agent even if agent_id is provided + + Raises: + ValueError: If agent or thread creation fails + """ + if self._agent is None or create_new_agent: + if self._agent_id and create_new_agent is False: + self._agent = await self._project_client.agents.get_agent(agent_id=self._agent_id) + else: + self._agent = await self._project_client.agents.create_agent( + name=self.name, + model=self._deployment_name, + description=self.description, + instructions=self._instructions, + tools=self._api_tools, + metadata=self._metadata, + response_format=self._response_format if self._response_format else None, # type: ignore + temperature=self._temperature, + tool_resources=self._tool_resources if self._tool_resources else None, # type: ignore + top_p=self._top_p, + ) + + if self._thread is None or create_new_thread: + if self._init_thread_id and create_new_thread is False: + self._thread = await self._project_client.agents.threads.get(thread_id=self._init_thread_id) + # Retrieve initial state only once + if not self._initial_state_retrieved: + await self._retrieve_initial_state() + self._initial_state_retrieved = True + else: + self._thread = await self._project_client.agents.threads.create() + + async def _retrieve_initial_state(self) -> None: + """ + Retrieve and store the initial state of messages in the thread. + + This method retrieves all message IDs from an existing thread to track which + messages were present before this agent instance started interacting with the thread. + It handles pagination to ensure all messages are captured. + """ + # Retrieve all initial message IDs + initial_message_ids: Set[str] = set() + async for msg in self._project_client.agents.messages.list( + thread_id=self.thread_id, + order=ListSortOrder.ASCENDING, + limit=100, + ): + initial_message_ids.add(msg.id) + self._initial_message_ids = initial_message_ids + + async def _execute_tool_call(self, tool_call: FunctionCall, cancellation_token: CancellationToken) -> str: + """ + Execute a tool call requested by the Azure AI agent. + + Args: + tool_call (FunctionCall): The function call information including name and arguments + cancellation_token (CancellationToken): Token for cancellation handling + + Returns: + str: The string representation of the tool call result + + Raises: + ValueError: If the requested tool is not available or no tools are registered + """ + if not self._original_tools: + raise ValueError("No tools are available.") + tool = next((t for t in self._original_tools if t.name == tool_call.name), None) + if tool is None: + raise ValueError(f"The tool '{tool_call.name}' is not available.") + arguments = json.loads(tool_call.arguments) + result = await tool.run_json(arguments, cancellation_token, call_id=tool_call.id) + return tool.return_value_as_string(result) + + async def _upload_files( + self, + file_paths: str | Iterable[str], + purpose: str = "assistant", + polling_interval: float = 0.5, + cancellation_token: Optional[CancellationToken] = None, + ) -> List[str]: + """ + Upload files to the Azure AI Assistant API. + + This method handles uploading one or more files to be used by the agent + and tracks their IDs in the agent's state. + + Args: + file_paths (str | Iterable[str]): Path(s) to file(s) to upload + purpose (str): The purpose of the file, defaults to "assistant" + polling_interval (float): Time to sleep between polling for file status + cancellation_token (Optional[CancellationToken]): Token for cancellation handling + + Returns: + List[str]: List of file IDs for the uploaded files + + Raises: + ValueError: If file upload fails + """ + if cancellation_token is None: + cancellation_token = CancellationToken() + + await self._ensure_initialized() + + if isinstance(file_paths, str): + file_paths = [file_paths] + + file_ids: List[str] = [] + for file_path in file_paths: + file_name = os.path.basename(file_path) + + file: FileInfo = await cancellation_token.link_future( + asyncio.ensure_future( + self._project_client.agents.files.upload_and_poll( + file_path=file_path, purpose=purpose, polling_interval=polling_interval + ) + ) + ) + + if file.status != FileState.PROCESSED: + raise ValueError(f"File upload failed with status {file.status}") + + trace_logger.debug(f"File uploaded successfully: {file.id}, {file_name}") + + file_ids.append(file.id) + self._uploaded_file_ids.append(file.id) + + return file_ids + + # Public Methods + async def on_messages( + self, + messages: Sequence[BaseChatMessage], + cancellation_token: Optional[CancellationToken] = None, + message_limit: int = 1, + ) -> Response: + """ + Process incoming messages and return a response from the Azure AI agent. + + This method is the primary entry point for interaction with the agent. + It delegates to on_messages_stream and returns the final response. + + Args: + messages (Sequence[BaseChatMessage]): The messages to process + cancellation_token (CancellationToken): Token for cancellation handling + message_limit (int, optional): Maximum number of messages to retrieve from the thread + + Returns: + Response: The agent's response, including the chat message and any inner events + + Raises: + AssertionError: If the stream doesn't return a final result + """ + async for message in self.on_messages_stream( + messages=messages, cancellation_token=cancellation_token, message_limit=message_limit + ): + if isinstance(message, Response): + return message + raise AssertionError("The stream should have returned the final result.") + + async def on_messages_stream( + self, + messages: Sequence[BaseChatMessage], + cancellation_token: Optional[CancellationToken] = None, + message_limit: int = 1, + polling_interval: float = 0.5, + ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]: + """ + Process incoming messages and yield streaming responses from the Azure AI agent. + + This method handles the complete interaction flow with the Azure AI agent: + 1. Processing input messages + 2. Creating and monitoring a run + 3. Handling tool calls and their results + 4. Retrieving and returning the agent's final response + + The method yields events during processing (like tool calls) and finally yields + the complete Response with the agent's message. + + Args: + messages (Sequence[BaseChatMessage]): The messages to process + cancellation_token (CancellationToken): Token for cancellation handling + message_limit (int, optional): Maximum number of messages to retrieve from the thread + polling_interval (float, optional): Time to sleep between polling for run status + + Yields: + AgentEvent | ChatMessage | Response: Events during processing and the final response + + Raises: + ValueError: If the run fails or no message is received from the assistant + """ + if cancellation_token is None: + cancellation_token = CancellationToken() + + await self._ensure_initialized() + + # Process all messages in sequence + for message in messages: + if isinstance(message, (TextMessage, MultiModalMessage)): + await self.handle_text_message(str(message.content), cancellation_token) + elif isinstance(message, (StopMessage, HandoffMessage)): + await self.handle_text_message(message.content, cancellation_token) + + # Inner messages for tool calls + inner_messages: List[AgentEvent | ChatMessage] = [] + + # Create and start a run + run: ThreadRun = await cancellation_token.link_future( + asyncio.ensure_future( + self._project_client.agents.runs.create( + thread_id=self.thread_id, + agent_id=self._get_agent_id, + ) + ) + ) + + # Wait for run completion by polling + while True: + run = await cancellation_token.link_future( + asyncio.ensure_future( + self._project_client.agents.runs.get( + thread_id=self.thread_id, + run_id=run.id, + ) + ) + ) + + if run.status == RunStatus.FAILED: + raise ValueError(f"Run failed: {run.last_error}") + + # If the run requires action (function calls), execute tools and continue + if run.status == RunStatus.REQUIRES_ACTION and run.required_action is not None: + tool_calls: List[FunctionCall] = [] + submit_tool_outputs = getattr(run.required_action, "submit_tool_outputs", None) + if submit_tool_outputs and hasattr(submit_tool_outputs, "tool_calls"): + for required_tool_call in submit_tool_outputs.tool_calls: + if required_tool_call.type == "function": + tool_calls.append( + FunctionCall( + id=required_tool_call.id, + name=required_tool_call.function.name, + arguments=required_tool_call.function.arguments, + ) + ) + + # Add tool call message to inner messages + tool_call_msg = ToolCallRequestEvent(source=self.name, content=tool_calls) + inner_messages.append(tool_call_msg) + trace_logger.debug(tool_call_msg) + yield tool_call_msg + + # Execute tool calls and get results + tool_outputs: List[FunctionExecutionResult] = [] + + # TODO: Support parallel execution of tool calls + + for tool_call in tool_calls: + try: + result = await self._execute_tool_call(tool_call, cancellation_token) + is_error = False + except Exception as e: + result = f"Error: {e}" + is_error = True + tool_outputs.append( + FunctionExecutionResult( + content=result, call_id=tool_call.id, is_error=is_error, name=tool_call.name + ) + ) + + # Add tool result message to inner messages + tool_result_msg = ToolCallExecutionEvent(source=self.name, content=tool_outputs) + inner_messages.append(tool_result_msg) + trace_logger.debug(tool_result_msg) + yield tool_result_msg + + # Submit tool outputs back to the run + run = await cancellation_token.link_future( + asyncio.ensure_future( + self._project_client.agents.runs.submit_tool_outputs( + thread_id=self.thread_id, + run_id=run.id, + tool_outputs=[ToolOutput(tool_call_id=t.call_id, output=t.content) for t in tool_outputs], + ) + ) + ) + continue + + if run.status == RunStatus.COMPLETED: + break + + # TODO support for parameter to control polling interval + await asyncio.sleep(polling_interval) + + # After run is completed, get the messages + trace_logger.debug("Retrieving messages from thread") + # Collect up to message_limit messages in DESCENDING order, support cancellation + agent_messages: List[ThreadMessage] = [] + async for msg in self._project_client.agents.messages.list( + thread_id=self.thread_id, + order=ListSortOrder.DESCENDING, + limit=message_limit, + ): + if cancellation_token.is_cancelled(): + trace_logger.debug("Message retrieval cancelled by token.") + break + agent_messages.append(msg) + if len(agent_messages) >= message_limit: + break + if not agent_messages: + raise ValueError("No messages received from assistant") + + # Get the last message from the agent (role=AGENT) + last_message: Optional[ThreadMessage] = next( + (m for m in agent_messages if getattr(m, "role", None) == "agent"), None + ) + if not last_message: + trace_logger.debug("No message with AGENT role found, falling back to first message") + last_message = agent_messages[0] # Fallback to first message + if not getattr(last_message, "content", None): + raise ValueError("No content in the last message") + + # Extract text content + message_text = "" + for text_message in last_message.text_messages: + message_text += text_message.text.value + + # Extract citations + citations: list[Any] = [] + + # Try accessing annotations directly + + annotations = getattr(last_message, "annotations", []) + + if isinstance(annotations, list) and annotations: + annotations = cast(List[MessageTextUrlCitationAnnotation], annotations) + + trace_logger.debug(f"Found {len(annotations)} annotations") + for annotation in annotations: + if hasattr(annotation, "url_citation"): # type: ignore + trace_logger.debug(f"Citation found: {annotation.url_citation.url}") + citations.append( + {"url": annotation.url_citation.url, "title": annotation.url_citation.title, "text": None} # type: ignore + ) + # For backwards compatibility + elif hasattr(last_message, "url_citation_annotations") and last_message.url_citation_annotations: + url_annotations = cast(List[Any], last_message.url_citation_annotations) + + trace_logger.debug(f"Found {len(url_annotations)} URL citations") + + for annotation in url_annotations: + citations.append( + {"url": annotation.url_citation.url, "title": annotation.url_citation.title, "text": None} # type: ignore + ) + + elif hasattr(last_message, "file_citation_annotations") and last_message.file_citation_annotations: + file_annotations = cast(List[Any], last_message.file_citation_annotations) + + trace_logger.debug(f"Found {len(file_annotations)} URL citations") + + for annotation in file_annotations: + citations.append( + {"file_id": annotation.file_citation.file_id, "title": None, "text": annotation.file_citation.quote} # type: ignore + ) + + trace_logger.debug(f"Total citations extracted: {len(citations)}") + + # Create the response message with citations as JSON string + chat_message = TextMessage( + source=self.name, content=message_text, metadata={"citations": json.dumps(citations)} if citations else {} + ) + + # Return the assistant's response as a Response with inner messages + yield Response(chat_message=chat_message, inner_messages=inner_messages) + + async def handle_text_message(self, content: str, cancellation_token: Optional[CancellationToken] = None) -> None: + """ + Handle a text message by adding it to the conversation thread. + + Args: + content (str): The text content of the message + cancellation_token (CancellationToken): Token for cancellation handling + + Returns: + None + """ + + if cancellation_token is None: + cancellation_token = CancellationToken() + + await cancellation_token.link_future( + asyncio.ensure_future( + self._project_client.agents.messages.create( + thread_id=self.thread_id, + role=MessageRole.USER, + content=content, + ) + ) + ) + + async def on_reset(self, cancellation_token: CancellationToken) -> None: + """ + Reset the agent's conversation by creating a new thread. + + This method allows for resetting a conversation without losing the agent + definition or capabilities. It creates a new thread for fresh conversations. + + Note: Currently the Azure AI Agent API has no support for deleting messages, + so a new thread is created instead. + + Args: + cancellation_token (CancellationToken): Token for cancellation handling + """ + # This will enforce the creation of a new thread + await self._ensure_initialized(create_new_thread=True) + + async def save_state(self) -> Mapping[str, Any]: + """ + Save the current state of the agent for future restoration. + + This method serializes the agent's state including IDs for the agent, thread, + messages, and associated resources like vector stores and uploaded files. + + Returns: + Mapping[str, Any]: A dictionary containing the serialized state data + """ + state = AzureAIAgentState( + agent_id=self._agent.id if self._agent else self._agent_id, + thread_id=self._thread.id if self._thread else self._init_thread_id, + initial_message_ids=list(self._initial_message_ids), + vector_store_id=self._vector_store_id, + uploaded_file_ids=self._uploaded_file_ids, + ) + return state.model_dump() + + async def load_state(self, state: Mapping[str, Any]) -> None: + """ + Load a previously saved state into this agent. + + This method deserializes and restores a previously saved agent state, + setting up the agent to continue a previous conversation or session. + + Args: + state (Mapping[str, Any]): The previously saved state dictionary + """ + agent_state = AzureAIAgentState.model_validate(state) + self._agent_id = agent_state.agent_id + self._init_thread_id = agent_state.thread_id + self._initial_message_ids = set(agent_state.initial_message_ids) + self._vector_store_id = agent_state.vector_store_id + self._uploaded_file_ids = agent_state.uploaded_file_ids + + async def on_upload_for_code_interpreter( + self, + file_paths: str | Iterable[str], + cancellation_token: Optional[CancellationToken] = None, + polling_interval: float = 0.5, + ) -> None: + """ + Upload files to be used with the code interpreter tool. + + This method uploads files for the agent's code interpreter tool and + updates the thread's tool resources to include these files. + + Args: + file_paths (str | Iterable[str]): Path(s) to file(s) to upload + cancellation_token (Optional[CancellationToken]): Token for cancellation handling + polling_interval (float): Time to sleep between polling for file status + + Raises: + ValueError: If file upload fails or the agent doesn't have code interpreter capability + """ + if cancellation_token is None: + cancellation_token = CancellationToken() + + await self._ensure_initialized() + + file_ids = await self._upload_files( + file_paths=file_paths, + cancellation_token=cancellation_token, + polling_interval=polling_interval, + purpose=FilePurpose.AGENTS, + ) + + # Update thread with the new files + thread: AgentThread = await cancellation_token.link_future( + asyncio.ensure_future(self._project_client.agents.threads.get(thread_id=self.thread_id)) + ) + + tool_resources: ToolResources = thread.tool_resources or ToolResources() + code_interpreter_resource = tool_resources.code_interpreter or CodeInterpreterToolResource() + existing_file_ids: List[str] = code_interpreter_resource.file_ids or [] + existing_file_ids.extend(file_ids) + + await cancellation_token.link_future( + asyncio.ensure_future( + self._project_client.agents.threads.update( + thread_id=self.thread_id, + tool_resources=ToolResources( + code_interpreter=CodeInterpreterToolResource(file_ids=existing_file_ids) + ), + ) + ) + ) + + async def on_upload_for_file_search( + self, + file_paths: str | Iterable[str], + cancellation_token: CancellationToken, + vector_store_name: Optional[str] = None, + data_sources: Optional[List[VectorStoreDataSource]] = None, + expires_after: Optional[VectorStoreExpirationPolicy] = None, + chunking_strategy: Optional[VectorStoreChunkingStrategyRequest] = None, + vector_store_metadata: Optional[Dict[str, str]] = None, + vector_store_polling_interval: float = 1, + ) -> None: + """ + Upload files to be used with the file search tool. + + This method handles uploading files for the file search capability, creating a vector + store if necessary, and updating the agent's configuration to use the vector store. + + Args: + file_paths (str | Iterable[str]): Path(s) to file(s) to upload + cancellation_token (CancellationToken): Token for cancellation handling + vector_store_name (Optional[str]): Name to assign to the vector store if creating a new one + data_sources (Optional[List[VectorStoreDataSource]]): Additional data sources for the vector store + expires_after (Optional[VectorStoreExpirationPolicy]): Expiration policy for vector store content + chunking_strategy (Optional[VectorStoreChunkingStrategyRequest]): Strategy for chunking file content + vector_store_metadata (Optional[Dict[str, str]]): Additional metadata for the vector store + vector_store_polling_interval (float): Time to sleep between polling for vector store status + + Raises: + ValueError: If file search is not enabled for this agent or file upload fails + """ + await self._ensure_initialized() + + # Check if file_search is enabled in tools + if not any(isinstance(tool, FileSearchToolDefinition) for tool in self._api_tools): + raise ValueError( + "File search is not enabled for this assistant. Add a file_search tool when creating the assistant." + ) + + # Create vector store if not already created + if self._vector_store_id is None: + vector_store: VectorStore = await cancellation_token.link_future( + asyncio.ensure_future( + self._project_client.agents.vector_stores.create_and_poll( + file_ids=[], + name=vector_store_name, + data_sources=data_sources, + expires_after=expires_after, + chunking_strategy=chunking_strategy, + metadata=vector_store_metadata, + polling_interval=vector_store_polling_interval, + ) + ) + ) + self._vector_store_id = vector_store.id + + # Update assistant with vector store ID + await cancellation_token.link_future( + asyncio.ensure_future( + self._project_client.agents.update_agent( + agent_id=self._get_agent_id, + tools=self._api_tools, + tool_resources=ToolResources( + file_search=FileSearchToolResource(vector_store_ids=[self._vector_store_id]) + ), + ) + ) + ) + + file_ids = await self._upload_files( + file_paths=file_paths, cancellation_token=cancellation_token, purpose=FilePurpose.AGENTS + ) + + # Create file batch with the file IDs + await cancellation_token.link_future( + asyncio.ensure_future( + self._project_client.agents.vector_store_file_batches.create_and_poll( + vector_store_id=self._vector_store_id, + file_ids=file_ids, + polling_interval=vector_store_polling_interval, + ) + ) + ) + + async def close(self) -> None: + """ + Close the Azure AI agent and release any resources. + """ + await self._project_client.close() + + +if __name__ == "__main__": + # Example usage of AzureAIAgent + # Replace with your actual endpoint and credentials + """ + TODO: + [X] Support for file upload + [] Support for sharepoint grounding + [] Support for azure function grounding + [X] Support for file search + [X] Support for custom function calling + [X] Add metadata to the thread (agent_id, source ="AUTODGEN_AGENT") + """ diff --git a/python/packages/autogen-ext/src/autogen_ext/agents/azure/_types.py b/python/packages/autogen-ext/src/autogen_ext/agents/azure/_types.py new file mode 100644 index 000000000000..e431d4265dfd --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/agents/azure/_types.py @@ -0,0 +1,61 @@ +from typing import Any, Awaitable, Callable, Iterable, List, Literal, Optional, TypeGuard, Union + +from autogen_core.tools import Tool +from pydantic import BaseModel, Field + +from azure.ai.agents.models import ( + AzureAISearchToolDefinition, + AzureFunctionToolDefinition, + BingGroundingToolDefinition, + CodeInterpreterToolDefinition, + FileSearchToolDefinition, + MessageTextUrlCitationAnnotation, +) + +ListToolType = Iterable[ + Union[ + Literal[ + "file_search", + "code_interpreter", + "bing_grounding", + "azure_ai_search", + "azure_function", + ], + BingGroundingToolDefinition, + CodeInterpreterToolDefinition, + AzureAISearchToolDefinition, + FileSearchToolDefinition, + AzureFunctionToolDefinition, + Tool, + Callable[..., Any], + Callable[..., Awaitable[Any]], + ] +] + + +class AzureAIAgentState(BaseModel): + """ + Represents the state of an AzureAIAgent that can be saved and loaded. + + This state model keeps track of persistent information about an agent session + including agent and thread identifiers, message history, and associated resources. + + Attributes: + type (str): The type identifier for the state object, always "AzureAIAgentState" + agent_id (Optional[str]): The ID of the Azure AI agent + thread_id (Optional[str]): The ID of the conversation thread + initial_message_ids (List[str]): List of message IDs from the initial state + vector_store_id (Optional[str]): The ID of the associated vector store for file search + uploaded_file_ids (List[str]): List of IDs for files uploaded to the agent + """ + + type: str = Field(default="AzureAIAgentState") + agent_id: Optional[str] = None + thread_id: Optional[str] = None + initial_message_ids: List[str] = Field(default_factory=list) + vector_store_id: Optional[str] = None + uploaded_file_ids: List[str] = Field(default_factory=list) + + +def has_annotations(obj: Any) -> TypeGuard[list[MessageTextUrlCitationAnnotation]]: + return obj is not None and isinstance(obj, list) diff --git a/python/packages/autogen-ext/src/autogen_ext/agents/file_surfer/_file_surfer.py b/python/packages/autogen-ext/src/autogen_ext/agents/file_surfer/_file_surfer.py index aec34cc6364b..91cd017204e2 100644 --- a/python/packages/autogen-ext/src/autogen_ext/agents/file_surfer/_file_surfer.py +++ b/python/packages/autogen-ext/src/autogen_ext/agents/file_surfer/_file_surfer.py @@ -6,8 +6,7 @@ from autogen_agentchat.agents import BaseChatAgent from autogen_agentchat.base import Response from autogen_agentchat.messages import ( - ChatMessage, - MultiModalMessage, + BaseChatMessage, TextMessage, ) from autogen_agentchat.utils import remove_images @@ -85,16 +84,12 @@ def __init__( self._browser = MarkdownFileBrowser(viewport_size=1024 * 5, base_path=base_path) @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: return (TextMessage,) - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: for chat_message in messages: - if isinstance(chat_message, TextMessage | MultiModalMessage): - self._chat_history.append(UserMessage(content=chat_message.content, source=chat_message.source)) - else: - raise ValueError(f"Unexpected message in FileSurfer: {chat_message}") - + self._chat_history.append(chat_message.to_model_message()) try: _, content = await self._generate_reply(cancellation_token=cancellation_token) self._chat_history.append(AssistantMessage(content=content, source=self.name)) diff --git a/python/packages/autogen-ext/src/autogen_ext/agents/openai/__init__.py b/python/packages/autogen-ext/src/autogen_ext/agents/openai/__init__.py index c5460d5b443f..91936e62ef49 100644 --- a/python/packages/autogen-ext/src/autogen_ext/agents/openai/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/agents/openai/__init__.py @@ -1,10 +1,7 @@ -try: - from ._openai_assistant_agent import OpenAIAssistantAgent -except ImportError as e: - raise ImportError( - "Dependencies for OpenAIAssistantAgent not found. " - 'Please install autogen-ext with the "openai" extra: ' - 'pip install "autogen-ext[openai]"' - ) from e +from ._openai_agent import OpenAIAgent +from ._openai_assistant_agent import OpenAIAssistantAgent -__all__ = ["OpenAIAssistantAgent"] +__all__ = [ + "OpenAIAgent", + "OpenAIAssistantAgent", +] diff --git a/python/packages/autogen-ext/src/autogen_ext/agents/openai/_openai_agent.py b/python/packages/autogen-ext/src/autogen_ext/agents/openai/_openai_agent.py new file mode 100644 index 000000000000..db9360d4160c --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/agents/openai/_openai_agent.py @@ -0,0 +1,1093 @@ +import asyncio +import json +import logging +import warnings +from typing import ( + Any, + AsyncGenerator, + Dict, + Iterable, + List, + Literal, + Mapping, + Optional, + Sequence, + Type, + Union, + cast, +) + +from autogen_agentchat import EVENT_LOGGER_NAME +from autogen_agentchat.agents import BaseChatAgent +from autogen_agentchat.base import Response +from autogen_agentchat.messages import ( + AgentEvent, + BaseChatMessage, + ChatMessage, + HandoffMessage, + MultiModalMessage, + StopMessage, + TextMessage, + ToolCallSummaryMessage, +) +from autogen_core import CancellationToken, Component, ComponentModel, FunctionCall +from autogen_core.models import UserMessage +from autogen_core.tools import Tool +from pydantic import BaseModel, Field +from typing_extensions import NotRequired, TypedDict + +from openai import AsyncAzureOpenAI, AsyncOpenAI # type: ignore + +# Number of characters to display when previewing image content in logs and UI +# Base64 encoded images can be very long, so we truncate for readability +IMAGE_CONTENT_PREVIEW_LENGTH = 50 + +# NOTE: We use the new Responses API, so ChatCompletion imports are not needed. + +event_logger = logging.getLogger(EVENT_LOGGER_NAME) + + +# TypedDict classes for built-in tool configurations +class FileSearchToolConfig(TypedDict): + """Configuration for file_search tool.""" + + type: Literal["file_search"] + vector_store_ids: List[str] # required - The IDs of the vector stores to search + max_num_results: NotRequired[int] # optional + ranking_options: NotRequired[Dict[str, Any]] # optional + filters: NotRequired[Dict[str, Any]] # optional + + +class WebSearchToolConfig(TypedDict): + """Configuration for web_search_preview tool.""" + + type: Literal["web_search_preview"] + search_context_size: NotRequired[int] # optional + user_location: NotRequired[Union[str, Dict[str, Any]]] # optional - Can be string or structured location + + +class ComputerUseToolConfig(TypedDict): + """Configuration for computer_use_preview tool.""" + + type: Literal["computer_use_preview"] + display_height: int # required - Display height in pixels + display_width: int # required - Display width in pixels + environment: str # required - Environment type for computer use + + +class MCPToolConfig(TypedDict): + """Configuration for mcp tool.""" + + type: Literal["mcp"] + server_label: str # required - Label for the MCP server + server_url: str # required - URL of the MCP server + allowed_tools: NotRequired[List[str]] # optional - List of allowed tools + headers: NotRequired[Dict[str, str]] # optional - HTTP headers for requests + require_approval: NotRequired[bool] # optional - Whether to require user approval + + +class CodeInterpreterToolConfig(TypedDict): + """Configuration for code_interpreter tool.""" + + type: Literal["code_interpreter"] + container: str # required - Container configuration for code execution + + +class ImageGenerationToolConfig(TypedDict): + """Configuration for image_generation tool.""" + + type: Literal["image_generation"] + background: NotRequired[str] # optional - Background color or image + input_image_mask: NotRequired[str] # optional - Mask for input image editing + + +class LocalShellToolConfig(TypedDict): + """Configuration for local_shell tool. + + WARNING: This tool is only supported with the 'codex-mini-latest' model + and is available exclusively through the Responses API. + """ + + type: Literal["local_shell"] + # Note: local_shell currently has no additional parameters in the API + + +# Union type for all built-in tool configurations +BuiltinToolConfig = Union[ + FileSearchToolConfig, + WebSearchToolConfig, + ComputerUseToolConfig, + MCPToolConfig, + CodeInterpreterToolConfig, + ImageGenerationToolConfig, + LocalShellToolConfig, +] + + +# Define ImageMessage class early since it's used in _convert_message_to_openai_message +class ImageMessage(BaseChatMessage): + """A message containing an image.""" + + content: str # URL or base64 string + + def to_model_message(self) -> UserMessage: + return UserMessage(content=self.content, source=self.source) + + def to_model_text(self) -> str: + return "[image]" + + def to_text(self) -> str: + # Truncate long image content (especially base64) for better readability + # While still showing enough of the URL or content to be identifiable + if len(self.content) > IMAGE_CONTENT_PREVIEW_LENGTH: + return f"[Image: {self.content[:IMAGE_CONTENT_PREVIEW_LENGTH]}...]" + return f"[Image: {self.content}]" + + +def _convert_tool_to_function_schema(tool: Tool) -> Dict[str, Any]: + schema = tool.schema + parameters: Dict[str, object] = {} + if "parameters" in schema: + parameters = { + "type": schema["parameters"]["type"], + "properties": schema["parameters"]["properties"], + } + if "required" in schema["parameters"]: + parameters["required"] = schema["parameters"]["required"] + + return { + "name": schema["name"], + "description": schema.get("description", ""), + "parameters": parameters, + } + + +class OpenAIMessageContent(TypedDict): + type: str + text: str + + +class OpenAIImageUrlContent(TypedDict): + url: str + + +class OpenAIImageContent(TypedDict): + type: str + image_url: OpenAIImageUrlContent + + +class OpenAIMessage(TypedDict): + role: str + content: Union[str, List[Union[OpenAIMessageContent, OpenAIImageContent]]] + + +def _convert_message_to_openai_message( + message: Union[TextMessage, MultiModalMessage, StopMessage, ToolCallSummaryMessage, HandoffMessage], +) -> OpenAIMessage: + """Convert an AutoGen message to an OpenAI message format.""" + if isinstance(message, TextMessage): + if message.source == "user": + return {"role": "user", "content": str(message.content)} + elif message.source == "system": + return {"role": "system", "content": str(message.content)} + elif message.source == "assistant": + return {"role": "assistant", "content": str(message.content)} + else: + return {"role": "user", "content": str(message.content)} + elif isinstance(message, MultiModalMessage): + content_parts: List[Union[OpenAIMessageContent, OpenAIImageContent]] = [] + for part in message.content: + if isinstance(part, TextMessage): + content_parts.append({"type": "text", "text": str(part.content)}) + elif isinstance(part, ImageMessage): + image_content = str(part.content) + content_parts.append({"type": "image_url", "image_url": {"url": image_content}}) + return {"role": "user", "content": content_parts} + else: + return {"role": "user", "content": str(message.content)} + + +class OpenAIAgentState(BaseModel): + type: str = Field(default="OpenAIAgentState") + response_id: Optional[str] = None + history: List[Dict[str, Any]] = Field(default_factory=list) + + +# Union type for tool configurations in the config schema +ToolConfigUnion = Union[ComponentModel, BuiltinToolConfig, str] + + +class OpenAIAgentConfig(BaseModel): + """Configuration model for OpenAI agent that supports both custom tools and built-in tools. + + .. versionchanged:: v0.7.0 + Added support for built-in tools in JSON configuration via _to_config and _from_config methods. + The tools field now accepts ComponentModel (for custom tools), built-in tool configurations + (dict format), and built-in tool names (string format). + """ + + name: str + description: str + model: str + instructions: str + tools: List[ToolConfigUnion] | None = None + temperature: Optional[float] = 1 + max_output_tokens: Optional[int] = None + json_mode: bool = False + store: bool = True + truncation: str = "disabled" + + +class FunctionExecutionResult(BaseModel): + """Result of a function execution.""" + + content: str + call_id: str + name: str + is_error: bool = False + + +class OpenAIAgent(BaseChatAgent, Component[OpenAIAgentConfig]): + """ + An agent implementation that uses the OpenAI Responses API to generate responses. + + Installation: + + .. code-block:: bash + + pip install "autogen-ext[openai]" + # pip install "autogen-ext[openai,azure]" # For Azure OpenAI Assistant + + This agent leverages the Responses API to generate responses with capabilities like: + + * Custom function calling + * Multi-turn conversations + * Built-in tool support (file_search, code_interpreter, web_search_preview, etc.) + + .. versionchanged:: v0.7.0 + + Added support for built-in tool types like file_search, web_search_preview, + code_interpreter, computer_use_preview, image_generation, and mcp. + Added support for tool configurations with required and optional parameters. + + Built-in tools are split into two categories: + + **Tools that can use string format** (no required parameters): + + - web_search_preview: Can be used as "web_search_preview" or with optional config + (user_location, search_context_size) + - image_generation: Can be used as "image_generation" or with optional config (background, input_image_mask) + - local_shell: Can be used as "local_shell" (WARNING: Only works with codex-mini-latest model) + + **Tools that REQUIRE dict configuration** (have required parameters): + + - file_search: MUST use dict with vector_store_ids (List[str]) + - computer_use_preview: MUST use dict with display_height (int), display_width (int), environment (str) + - code_interpreter: MUST use dict with container (str) + - mcp: MUST use dict with server_label (str), server_url (str) + + Using required-parameter tools in string format will raise a ValueError with helpful error messages. + The tools parameter type annotation only accepts string values for tools that don't require parameters. + + + Args: + name (str): Name of the agent + description (str): Description of the agent's purpose + client (Union[AsyncOpenAI, AsyncAzureOpenAI]): OpenAI client instance + model (str): Model to use (e.g. "gpt-4.1") + instructions (str): System instructions for the agent + tools (Optional[Iterable[Union[str, BuiltinToolConfig, Tool]]]): Tools the agent can use. + Supported string values (no required parameters): "web_search_preview", "image_generation", "local_shell". + Dict values can provide configuration for built-in tools with parameters. + Required parameters for built-in tools: + - file_search: vector_store_ids (List[str]) + - computer_use_preview: display_height (int), display_width (int), environment (str) + - code_interpreter: container (str) + - mcp: server_label (str), server_url (str) + Optional parameters for built-in tools: + - file_search: max_num_results (int), ranking_options (dict), filters (dict) + - web_search_preview: user_location (str or dict), search_context_size (int) + - image_generation: background (str), input_image_mask (str) + - mcp: allowed_tools (List[str]), headers (dict), require_approval (bool) + Special tools with model restrictions: + - local_shell: Only works with "codex-mini-latest" model (WARNING: Very limited support) + Also accepts custom Tool objects for function calling. + temperature (Optional[float]): Temperature for response generation (default: 1) + max_output_tokens (Optional[int]): Maximum output tokens + json_mode (bool): Whether to use JSON mode (default: False) + store (bool): Whether to store conversations (default: True) + truncation (str): Truncation strategy (default: "disabled") + + Example: + + Basic usage with built-in tools: + + .. code-block:: python + + from openai import AsyncOpenAI + from autogen_core import CancellationToken + from autogen_ext.agents.openai import OpenAIAgent + from autogen_agentchat.messages import TextMessage + import logging + + + async def example(): + cancellation_token = CancellationToken() + client = AsyncOpenAI() + agent = OpenAIAgent( + name="Simple Agent", + description="A simple OpenAI agent using the Responses API", + client=client, + model="gpt-4o", + instructions="You are a helpful assistant.", + tools=["web_search_preview", "image_generation"], # Only tools without required params + ) + response = await agent.on_messages( + [TextMessage(source="user", content="Search for recent AI developments")], cancellation_token + ) + logging.info(response) + + Usage with configured built-in tools: + + .. code-block:: python + + from openai import AsyncOpenAI + from autogen_core import CancellationToken + from autogen_ext.agents.openai import OpenAIAgent + from autogen_agentchat.messages import TextMessage + import logging + + + async def example_with_configs(): + cancellation_token = CancellationToken() + client = AsyncOpenAI() + + # Configure tools with required and optional parameters + tools = [ + { + "type": "file_search", + "vector_store_ids": ["vs_abc123"], # required + "max_num_results": 10, # optional + }, + { + "type": "computer_use_preview", + "display_height": 1024, # required + "display_width": 1280, # required + "environment": "desktop", # required + }, + { + "type": "code_interpreter", + "container": "python-3.11", # required + }, + { + "type": "mcp", + "server_label": "my-mcp-server", # required + "server_url": "http://localhost:3000", # required + }, + { + "type": "web_search_preview", + "user_location": { # optional - structured location + "type": "approximate", # required: "approximate" or "exact" + "country": "US", # optional + "region": "CA", # optional + "city": "San Francisco", # optional + }, + "search_context_size": 5, # optional + }, + "image_generation", # Simple tools can still use string format + ] + + agent = OpenAIAgent( + name="Configured Agent", + description="An agent with configured tools", + client=client, + model="gpt-4o", + instructions="You are a helpful assistant with specialized tools.", + tools=tools, # type: ignore + ) + response = await agent.on_messages( + [TextMessage(source="user", content="Search for recent AI developments")], cancellation_token + ) + logging.info(response) + + Mixed usage with custom function tools: + + .. code-block:: python + + import asyncio + import logging + from openai import AsyncOpenAI + from autogen_core import CancellationToken + from autogen_ext.agents.openai import OpenAIAgent + from autogen_agentchat.messages import TextMessage + from autogen_core.tools import FunctionTool + + + # Define a simple calculator function + async def calculate(a: int, b: int) -> int: + '''Simple function to add two numbers.''' + return a + b + + + # Wrap the calculate function as a tool + calculator = FunctionTool(calculate, description="A simple calculator tool") + + + async def example_mixed_tools(): + cancellation_token = CancellationToken() + client = AsyncOpenAI() + # Use the FunctionTool instance defined above + + agent = OpenAIAgent( + name="Mixed Tools Agent", + description="An agent with both built-in and custom tools", + client=client, + model="gpt-4o", + instructions="You are a helpful assistant with calculation and web search capabilities.", + tools=[ + "web_search_preview", + calculator, + {"type": "mcp", "server_label": "tools", "server_url": "http://localhost:3000"}, + ], + ) + response = await agent.on_messages( + [TextMessage(source="user", content="What's 2+2 and what's the weather like?")], + cancellation_token, + ) + logging.info(response) + + + asyncio.run(example_mixed_tools()) + + + """ + + component_config_schema = OpenAIAgentConfig + component_provider_override = "autogen_ext.agents.openai.OpenAIAgent" + + def __init__( + self: "OpenAIAgent", + name: str, + description: str, + client: Union[AsyncOpenAI, AsyncAzureOpenAI], + model: str, + instructions: str, + tools: Optional[ + Iterable[ + Union[ + Literal["web_search_preview", "image_generation", "local_shell"], + BuiltinToolConfig, + Tool, + ] + ] + ] = None, + temperature: Optional[float] = 1, + max_output_tokens: Optional[int] = None, + json_mode: bool = False, + store: bool = True, + truncation: str = "disabled", + ) -> None: + super().__init__(name, description) + self._client: Union[AsyncOpenAI, AsyncAzureOpenAI] = client + self._model: str = model + self._instructions: str = instructions + self._temperature: Optional[float] = temperature + self._max_output_tokens: Optional[int] = max_output_tokens + self._json_mode: bool = json_mode + self._store: bool = store + self._truncation: str = truncation + self._last_response_id: Optional[str] = None + self._message_history: List[Dict[str, Any]] = [] + self._tools: List[Dict[str, Any]] = [] + self._tool_map: Dict[str, Tool] = {} + if tools is not None: + for tool in tools: + if isinstance(tool, str): + # Handle built-in tool types + self._add_builtin_tool(tool) + elif isinstance(tool, dict) and "type" in tool: + # Handle configured built-in tools + self._add_configured_tool(tool) + elif isinstance(tool, Tool): + # Handle custom function tools + function_schema: Dict[str, Any] = { + "type": "function", + "function": _convert_tool_to_function_schema(tool), + } + self._tools.append(function_schema) + self._tool_map[tool.name] = tool + else: + raise ValueError(f"Unsupported tool type: {type(tool)}") + + def _add_builtin_tool(self, tool_name: str) -> None: + """Add a built-in tool by name.""" + # Skip if an identical tool has already been registered (idempotent behaviour) + if any(td.get("type") == tool_name for td in self._tools): + return # Duplicate – ignore rather than raise to stay backward-compatible + # Only allow string format for tools that don't require parameters + if tool_name == "web_search_preview": + self._tools.append({"type": "web_search_preview"}) + elif tool_name == "image_generation": + self._tools.append({"type": "image_generation"}) + elif tool_name == "local_shell": + # Special handling for local_shell - very limited model support + if self._model != "codex-mini-latest": + raise ValueError( + f"Tool 'local_shell' is only supported with model 'codex-mini-latest', " + f"but current model is '{self._model}'. " + f"This tool is available exclusively through the Responses API and has severe limitations. " + f"Consider using autogen_ext.tools.code_execution.PythonCodeExecutionTool with " + f"autogen_ext.code_executors.local.LocalCommandLineCodeExecutor for shell execution instead." + ) + self._tools.append({"type": "local_shell"}) + elif tool_name in ["file_search", "code_interpreter", "computer_use_preview", "mcp"]: + # These tools require specific parameters and must use dict configuration + raise ValueError( + f"Tool '{tool_name}' requires specific parameters and cannot be added using string format. " + f"Use dict configuration instead. Required parameters for {tool_name}: " + f"{self._get_required_params_help(tool_name)}" + ) + else: + raise ValueError(f"Unsupported built-in tool type: {tool_name}") + + def _get_required_params_help(self, tool_name: str) -> str: + """Get help text for required parameters of a tool.""" + help_text = { + "file_search": "vector_store_ids (List[str])", + "code_interpreter": "container (str)", + "computer_use_preview": "display_height (int), display_width (int), environment (str)", + "mcp": "server_label (str), server_url (str)", + } + return help_text.get(tool_name, "unknown parameters") + + def _add_configured_tool(self, tool_config: BuiltinToolConfig) -> None: + """Add a configured built-in tool with parameters.""" + tool_type = tool_config.get("type") + if not tool_type: + raise ValueError("Tool configuration must include 'type' field") + + # If an identical configuration is already present we simply ignore the new one (keeps API payload minimal) + if cast(Dict[str, Any], tool_config) in self._tools: + return + + # Initialize tool definition + tool_def: Dict[str, Any] = {} + + # Special validation for model-restricted tools + if tool_type == "local_shell": + if self._model != "codex-mini-latest": + raise ValueError( + f"Tool 'local_shell' is only supported with model 'codex-mini-latest', " + f"but current model is '{self._model}'. " + f"This tool is available exclusively through the Responses API and has severe limitations. " + f"Consider using autogen_ext.tools.code_execution.PythonCodeExecutionTool with " + f"autogen_ext.code_executors.local.LocalCommandLineCodeExecutor for shell execution instead." + ) + tool_def = {"type": "local_shell"} + + # For Responses API, built-in tools are defined directly without nesting + elif tool_type == "file_search": + # file_search requires vector_store_ids + fs_config = cast(FileSearchToolConfig, tool_config) + if "vector_store_ids" not in fs_config: + raise ValueError("file_search tool requires 'vector_store_ids' parameter") + + vector_store_ids = fs_config["vector_store_ids"] + if not isinstance(vector_store_ids, list) or not vector_store_ids: + raise ValueError("file_search 'vector_store_ids' must be a non-empty list of strings") + if not all(isinstance(vid, str) and vid.strip() for vid in vector_store_ids): + raise ValueError("file_search 'vector_store_ids' must contain non-empty strings") + + tool_def = {"type": "file_search", "vector_store_ids": vector_store_ids} + # Optional parameters + if "max_num_results" in fs_config: + max_results = fs_config["max_num_results"] + if not isinstance(max_results, int) or max_results <= 0: + raise ValueError("file_search 'max_num_results' must be a positive integer") + tool_def["max_num_results"] = max_results + if "ranking_options" in fs_config: + tool_def["ranking_options"] = fs_config["ranking_options"] + if "filters" in fs_config: + tool_def["filters"] = fs_config["filters"] + + elif tool_type == "web_search_preview": + # web_search_preview can have optional parameters + ws_config = cast(WebSearchToolConfig, tool_config) + tool_def = {"type": "web_search_preview"} + if "search_context_size" in ws_config: + context_size = ws_config["search_context_size"] + if not isinstance(context_size, int) or context_size <= 0: + raise ValueError("web_search_preview 'search_context_size' must be a positive integer") + tool_def["search_context_size"] = context_size + if "user_location" in ws_config: + user_location = ws_config["user_location"] + if isinstance(user_location, str): + if not user_location.strip(): + raise ValueError( + "web_search_preview 'user_location' must be a non-empty string when using string format" + ) + elif isinstance(user_location, dict): + if "type" not in user_location: + raise ValueError("web_search_preview 'user_location' dictionary must include 'type' field") + location_type = user_location["type"] + if location_type not in ["approximate", "exact"]: + raise ValueError("web_search_preview 'user_location' type must be 'approximate' or 'exact'") + # Optional fields: country, region, city can be validated if present + for optional_field in ["country", "region", "city"]: + if optional_field in user_location: + if ( + not isinstance(user_location[optional_field], str) + or not user_location[optional_field].strip() + ): + raise ValueError( + f"web_search_preview 'user_location' {optional_field} must be a non-empty string" + ) + else: + raise ValueError("web_search_preview 'user_location' must be a string or dictionary") + tool_def["user_location"] = user_location + + elif tool_type == "computer_use_preview": + # computer_use_preview requires display dimensions and environment + cu_config = cast(ComputerUseToolConfig, tool_config) + required_params = ["display_height", "display_width", "environment"] + for param in required_params: + if param not in cu_config: + raise ValueError(f"computer_use_preview tool requires '{param}' parameter") + + # Validate display dimensions + height = cu_config["display_height"] + width = cu_config["display_width"] + if not isinstance(height, int) or height <= 0: + raise ValueError("computer_use_preview 'display_height' must be a positive integer") + if not isinstance(width, int) or width <= 0: + raise ValueError("computer_use_preview 'display_width' must be a positive integer") + + # Validate environment + environment = cu_config["environment"] + if not isinstance(environment, str) or not environment.strip(): + raise ValueError("computer_use_preview 'environment' must be a non-empty string") + + tool_def = { + "type": "computer_use_preview", + "display_height": height, + "display_width": width, + "environment": environment, + } + + elif tool_type == "mcp": + # MCP requires server_label and server_url + mcp_config = cast(MCPToolConfig, tool_config) + required_params = ["server_label", "server_url"] + for param in required_params: + if param not in mcp_config: + raise ValueError(f"mcp tool requires '{param}' parameter") + + # Validate required parameters + server_label = mcp_config["server_label"] + server_url = mcp_config["server_url"] + if not isinstance(server_label, str) or not server_label.strip(): + raise ValueError("mcp 'server_label' must be a non-empty string") + if not isinstance(server_url, str) or not server_url.strip(): + raise ValueError("mcp 'server_url' must be a non-empty string") + + tool_def = {"type": "mcp", "server_label": server_label, "server_url": server_url} + # Optional parameters + if "allowed_tools" in mcp_config: + allowed_tools = mcp_config["allowed_tools"] + if not isinstance(allowed_tools, list): + raise ValueError("mcp 'allowed_tools' must be a list of strings") + if not all(isinstance(tool, str) for tool in allowed_tools): + raise ValueError("mcp 'allowed_tools' must contain only strings") + tool_def["allowed_tools"] = allowed_tools + if "headers" in mcp_config: + headers = mcp_config["headers"] + if not isinstance(headers, dict): + raise ValueError("mcp 'headers' must be a dictionary") + tool_def["headers"] = headers + if "require_approval" in mcp_config: + require_approval = mcp_config["require_approval"] + if not isinstance(require_approval, bool): + raise ValueError("mcp 'require_approval' must be a boolean") + tool_def["require_approval"] = require_approval + + elif tool_type == "code_interpreter": + # code_interpreter requires container + ci_config = cast(CodeInterpreterToolConfig, tool_config) + if "container" not in ci_config: + raise ValueError("code_interpreter tool requires 'container' parameter") + + container = ci_config["container"] + if not isinstance(container, str) or not container.strip(): + raise ValueError("code_interpreter 'container' must be a non-empty string") + + tool_def = {"type": "code_interpreter", "container": container} + + elif tool_type == "image_generation": + # image_generation can have optional parameters + ig_config = cast(ImageGenerationToolConfig, tool_config) + tool_def = {"type": "image_generation"} + if "background" in ig_config: + background = ig_config["background"] + if not isinstance(background, str) or not background.strip(): + raise ValueError("image_generation 'background' must be a non-empty string") + tool_def["background"] = background + if "input_image_mask" in ig_config: + input_image_mask = ig_config["input_image_mask"] + if not isinstance(input_image_mask, str) or not input_image_mask.strip(): + raise ValueError("image_generation 'input_image_mask' must be a non-empty string") + tool_def["input_image_mask"] = input_image_mask + + else: + raise ValueError(f"Unsupported built-in tool type: {tool_type}") + + self._tools.append(tool_def) + + def _convert_message_to_dict(self, message: OpenAIMessage) -> Dict[str, Any]: + """Convert an OpenAIMessage to a Dict[str, Any].""" + return dict(message) + + @property + def produced_message_types( + self: "OpenAIAgent", + ) -> Sequence[ + Union[ + Type[TextMessage], + Type[MultiModalMessage], + Type[StopMessage], + Type[ToolCallSummaryMessage], + Type[HandoffMessage], + ] + ]: + """Return the types of messages that this agent can produce.""" + return [TextMessage, MultiModalMessage, StopMessage, ToolCallSummaryMessage, HandoffMessage] + + async def _execute_tool_call( + self: "OpenAIAgent", tool_call: FunctionCall, cancellation_token: CancellationToken + ) -> FunctionExecutionResult: + tool_name = tool_call.name + if tool_name not in self._tool_map: + return FunctionExecutionResult( + content=f"Error: Tool '{tool_name}' is not available", + call_id=tool_call.id, + name=tool_name, + is_error=True, + ) + + tool = self._tool_map[tool_name] + try: + try: + arguments = json.loads(tool_call.arguments) + except json.JSONDecodeError as json_err: + return FunctionExecutionResult( + content=f"Error: Invalid JSON in tool arguments - {str(json_err)}", + call_id=tool_call.id, + name=tool_name, + is_error=True, + ) + + result = await tool.run_json(arguments, cancellation_token, call_id=tool_call.id) + return FunctionExecutionResult( + content=tool.return_value_as_string(result), call_id=tool_call.id, name=tool_name, is_error=False + ) + except Exception as e: + error_msg = f"Error: {str(e)}" + event_logger.warning(f"Tool execution error in {tool_name}: {error_msg}") + return FunctionExecutionResult(content=error_msg, call_id=tool_call.id, name=tool_name, is_error=True) + + def _build_api_parameters(self: "OpenAIAgent", messages: List[Dict[str, Any]]) -> Dict[str, Any]: + has_system_message = any(msg.get("role") == "system" for msg in messages) + if self._instructions and not has_system_message: + messages = [{"role": "system", "content": self._instructions}] + messages + api_params: Dict[str, Any] = { + "model": self._model, + "input": messages, # Responses API expects 'input' + } + if self._temperature is not None: + api_params["temperature"] = self._temperature + if self._max_output_tokens is not None: + api_params["max_output_tokens"] = self._max_output_tokens + if self._tools: + api_params["tools"] = self._tools + if self._json_mode: + api_params["text"] = {"type": "json_object"} + api_params["store"] = self._store + api_params["truncation"] = self._truncation + if self._last_response_id: + api_params["previous_response_id"] = self._last_response_id + return api_params + + async def on_messages( + self: "OpenAIAgent", messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken + ) -> Response: + response = None + inner_messages: List[ + Union[AgentEvent, TextMessage, MultiModalMessage, StopMessage, ToolCallSummaryMessage, HandoffMessage] + ] = [] + + async for msg in self.on_messages_stream(messages, cancellation_token): + if isinstance(msg, Response): + response = msg + # ModelClientStreamingChunkEvent does not exist in this version, so skip this check + else: + inner_messages.append(msg) + + if response is None: + raise ValueError("No response was generated") + + if response.inner_messages is None: + response.inner_messages = [] + + for msg in inner_messages: + if msg not in response.inner_messages: + response.inner_messages = list(response.inner_messages) + [msg] + + return response + + async def on_messages_stream( + self: "OpenAIAgent", messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken + ) -> AsyncGenerator[ + Union[ + AgentEvent, TextMessage, MultiModalMessage, StopMessage, ToolCallSummaryMessage, HandoffMessage, Response + ], + None, + ]: + input_messages: List[Dict[str, Any]] = [] + + if self._message_history: + input_messages.extend(self._message_history) + + for message in messages: + if isinstance( + message, (TextMessage, MultiModalMessage, StopMessage, ToolCallSummaryMessage, HandoffMessage) + ): + openai_message = _convert_message_to_openai_message(message) + dict_message = self._convert_message_to_dict(openai_message) + input_messages.append(dict_message) + self._message_history.append(dict_message) + else: + msg_content = str(cast(Any, message).content) if hasattr(message, "content") else str(message) + dict_message = {"role": "user", "content": msg_content} + input_messages.append(dict_message) + self._message_history.append(dict_message) + + inner_messages: List[AgentEvent | ChatMessage] = [] + + api_params = self._build_api_parameters(input_messages) + + try: + client = cast(Any, self._client) + response_obj = await cancellation_token.link_future( + asyncio.ensure_future(client.responses.create(**api_params)) + ) + content = getattr(response_obj, "output_text", None) + response_id = getattr(response_obj, "id", None) + self._last_response_id = response_id + # Use a readable placeholder when the API returns no content to aid debugging + content_str: str = str(content) if content is not None else "[no content returned]" + self._message_history.append({"role": "assistant", "content": content_str}) + final_message = TextMessage(source=self.name, content=content_str) + response = Response(chat_message=final_message, inner_messages=inner_messages) + yield response + except Exception as e: + error_message = f"Error generating response: {str(e)}" + event_logger.error(f"API error: {error_message}", exc_info=True) + error_response = TextMessage(source=self.name, content=error_message) + yield Response(chat_message=error_response, inner_messages=inner_messages) + + async def on_reset(self: "OpenAIAgent", cancellation_token: CancellationToken) -> None: + self._last_response_id = None + self._message_history = [] + + async def save_state(self: "OpenAIAgent") -> Mapping[str, Any]: + state = OpenAIAgentState( + response_id=self._last_response_id, + history=self._message_history, + ) + return state.model_dump() + + async def load_state(self: "OpenAIAgent", state: Mapping[str, Any]) -> None: + agent_state = OpenAIAgentState.model_validate(state) + self._last_response_id = agent_state.response_id + self._message_history = agent_state.history + + def _to_config(self: "OpenAIAgent") -> OpenAIAgentConfig: + """Convert the OpenAI agent to a declarative config. + + Serializes both custom Tool objects and built-in tools to their appropriate + configuration formats for JSON serialization. + + .. versionchanged:: v0.6.2 + Added support for serializing built-in tools alongside custom tools. + + Returns: + OpenAIAgentConfig: The configuration that can recreate this agent. + """ + # Serialize tools in the **original order** they were registered. We iterate over the + # internal ``self._tools`` list which contains both built-in tool definitions **and** the + # synthetic "function" records for custom :class:`Tool` objects. For the latter we + # convert the synthetic record back to a :class:`ComponentModel` by looking up the actual + # tool instance in ``self._tool_map``. This approach keeps ordering stable while still + # supporting full round-trip serialisation. + tool_configs: List[ToolConfigUnion] = [] + + for tool_def in self._tools: + # 1. Custom function tools are stored internally as ``{"type": "function", "function": {...}}``. + if tool_def.get("type") == "function": + fn_schema = cast(Dict[str, Any], tool_def.get("function", {})) + tool_name = fn_schema.get("name") # type: ignore[arg-type] + if tool_name and tool_name in self._tool_map: + tool_obj = self._tool_map[tool_name] + try: + if hasattr(tool_obj, "dump_component"): + component_model = cast(Any, tool_obj).dump_component() + tool_configs.append(component_model) + else: + component_model = ComponentModel( + provider="autogen_core.tools.FunctionTool", + component_type=None, + config={ + "name": tool_obj.name, + "description": getattr(tool_obj, "description", ""), + }, + ) + tool_configs.append(component_model) + except Exception as e: # pragma: no cover – extremely unlikely + warnings.warn( + f"Error serializing tool '{tool_name}': {e}", + stacklevel=2, + ) + component_model = ComponentModel( + provider="autogen_core.tools.FunctionTool", + component_type=None, + config={ + "name": tool_name or "unknown_tool", + "description": getattr(tool_obj, "description", ""), + }, + ) + tool_configs.append(component_model) + # 2. Built-in tools are already in their correct dict form – append verbatim. + elif "type" in tool_def: # built-in tool + tool_configs.append(cast(BuiltinToolConfig, tool_def)) + else: # pragma: no cover – should never happen + warnings.warn( + f"Encountered unexpected tool definition during serialisation: {tool_def}", + stacklevel=2, + ) + + return OpenAIAgentConfig( + name=self.name, + description=self.description, + model=self._model, + instructions=self._instructions, + tools=tool_configs if tool_configs else None, + temperature=self._temperature, + max_output_tokens=self._max_output_tokens, + json_mode=self._json_mode, + store=self._store, + truncation=self._truncation, + ) + + @classmethod + def _from_config(cls: Type["OpenAIAgent"], config: OpenAIAgentConfig) -> "OpenAIAgent": + """Create an OpenAI agent from a declarative config. + + Handles both custom Tool objects (from ComponentModel) and built-in tools + (from string or dict configurations). + + .. versionchanged:: v0.6.2 + Added support for loading built-in tools alongside custom tools. + + Args: + config: The configuration to load the agent from. + + Returns: + OpenAIAgent: The reconstructed agent. + """ + from openai import AsyncOpenAI + + client = AsyncOpenAI() + + tools: Optional[List[Union[str, BuiltinToolConfig, Tool]]] = None + if config.tools: + tools_list: List[Union[str, BuiltinToolConfig, Tool]] = [] + for tool_config in config.tools: + # Handle ComponentModel (custom Tool objects) + if isinstance(tool_config, ComponentModel): + try: + provider = tool_config.provider + module_name, class_name = provider.rsplit(".", 1) + module = __import__(module_name, fromlist=[class_name]) + tool_cls = getattr(module, class_name) + tool = tool_cls(**tool_config.config) + tools_list.append(cast(Tool, tool)) + except Exception as e: + warnings.warn(f"Error loading custom tool: {e}", stacklevel=2) + from autogen_core.tools import FunctionTool + + async def dummy_func(*args: Any, **kwargs: Any) -> str: + return "Tool not fully restored" + + tool = FunctionTool( + name=tool_config.config.get("name", "unknown_tool"), + description=tool_config.config.get("description", ""), + func=dummy_func, + ) + tools_list.append(tool) + + # Handle string format built-in tools + elif isinstance(tool_config, str): + tools_list.append(tool_config) + + # Handle dict format built-in tools + elif isinstance(tool_config, dict) and "type" in tool_config: + tools_list.append(tool_config) # type: ignore[arg-type] + + else: + warnings.warn(f"Unknown tool configuration format: {type(tool_config)}", stacklevel=2) + + tools = tools_list if tools_list else None + + return cls( + name=config.name, + description=config.description, + client=client, + model=config.model, + instructions=config.instructions, + tools=cast( + Optional[ + Iterable[ + Union[ + BuiltinToolConfig, + Tool, + Literal["web_search_preview", "image_generation", "local_shell"], + ] + ] + ], + tools, + ), + temperature=config.temperature, + max_output_tokens=config.max_output_tokens, + json_mode=config.json_mode, + store=config.store, + truncation=config.truncation, + ) + + # Add public API wrappers for configuration and tools + def to_config(self) -> OpenAIAgentConfig: + """Public wrapper for the private _to_config method.""" + return self._to_config() + + @classmethod + def from_config(cls, config: OpenAIAgentConfig) -> "OpenAIAgent": + """Public wrapper for the private _from_config classmethod.""" + return cls._from_config(config) + + @property + def tools(self) -> list[Any]: + """Public access to the agent's tools.""" + return self._tools + + @property + def model(self) -> str: + """Public access to the agent's model.""" + return self._model diff --git a/python/packages/autogen-ext/src/autogen_ext/agents/openai/_openai_assistant_agent.py b/python/packages/autogen-ext/src/autogen_ext/agents/openai/_openai_assistant_agent.py index 81b881704f77..38b27f248d32 100644 --- a/python/packages/autogen-ext/src/autogen_ext/agents/openai/_openai_assistant_agent.py +++ b/python/packages/autogen-ext/src/autogen_ext/agents/openai/_openai_assistant_agent.py @@ -24,18 +24,14 @@ from autogen_agentchat.agents import BaseChatAgent from autogen_agentchat.base import Response from autogen_agentchat.messages import ( - AgentEvent, - ChatMessage, - HandoffMessage, - MultiModalMessage, - StopMessage, + BaseAgentEvent, + BaseChatMessage, TextMessage, ToolCallExecutionEvent, ToolCallRequestEvent, ) -from autogen_core import CancellationToken, FunctionCall -from autogen_core.models._model_client import ChatCompletionClient -from autogen_core.models._types import FunctionExecutionResult +from autogen_core import CancellationToken, FunctionCall, Image +from autogen_core.models import ChatCompletionClient, FunctionExecutionResult from autogen_core.tools import FunctionTool, Tool from pydantic import BaseModel, Field @@ -52,6 +48,12 @@ from openai.types.beta.function_tool_param import FunctionToolParam from openai.types.beta.thread import Thread, ToolResources, ToolResourcesCodeInterpreter from openai.types.beta.threads import Message, MessageDeleted, Run +from openai.types.beta.threads.image_url_content_block_param import ImageURLContentBlockParam +from openai.types.beta.threads.image_url_param import ImageURLParam +from openai.types.beta.threads.message_content_part_param import ( + MessageContentPartParam, +) +from openai.types.beta.threads.text_content_block_param import TextContentBlockParam from openai.types.shared_params.function_definition import FunctionDefinition from openai.types.vector_store import VectorStore @@ -95,7 +97,7 @@ class OpenAIAssistantAgent(BaseChatAgent): .. code-block:: bash - pip install "autogen-ext[openai]" + pip install "autogen-ext[openai]" # For OpenAI Assistant # pip install "autogen-ext[openai,azure]" # For Azure OpenAI Assistant @@ -144,7 +146,7 @@ async def example(): # Create an assistant with code interpreter assistant = OpenAIAssistantAgent( - name="Python Helper", + name="PythonHelper", description="Helps with Python programming", client=client, model="gpt-4", @@ -195,7 +197,7 @@ async def example(): # Create an assistant with code interpreter assistant = OpenAIAssistantAgent( - name="Python Helper", + name="PythonHelper", description="Helps with Python programming", client=client, model="gpt-4o", @@ -309,9 +311,9 @@ async def _ensure_initialized(self) -> None: """Ensure assistant and thread are created.""" if self._assistant is None: if self._assistant_id: - self._assistant = await self._client.beta.assistants.retrieve(assistant_id=self._assistant_id) + self._assistant = await self._client.beta.assistants.retrieve(assistant_id=self._assistant_id) # type: ignore[reportDeprecated] else: - self._assistant = await self._client.beta.assistants.create( + self._assistant = await self._client.beta.assistants.create( # type: ignore[reportDeprecated] model=self._model, description=self.description, instructions=self._instructions, @@ -325,9 +327,9 @@ async def _ensure_initialized(self) -> None: if self._thread is None: if self._init_thread_id: - self._thread = await self._client.beta.threads.retrieve(thread_id=self._init_thread_id) + self._thread = await self._client.beta.threads.retrieve(thread_id=self._init_thread_id) # type: ignore[reportDeprecated] else: - self._thread = await self._client.beta.threads.create() + self._thread = await self._client.beta.threads.create() # type: ignore[reportDeprecated] # Retrieve initial state only once if not self._initial_state_retrieved: @@ -340,7 +342,7 @@ async def _retrieve_initial_state(self) -> None: initial_message_ids: Set[str] = set() after: str | NotGiven = NOT_GIVEN while True: - msgs: AsyncCursorPage[Message] = await self._client.beta.threads.messages.list( + msgs: AsyncCursorPage[Message] = await self._client.beta.threads.messages.list( # type: ignore[reportDeprecated] self._thread_id, after=after, order="asc", limit=100 ) for msg in msgs.data: @@ -351,7 +353,7 @@ async def _retrieve_initial_state(self) -> None: self._initial_message_ids = initial_message_ids @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: """The types of messages that the assistant agent produces.""" return (TextMessage,) @@ -387,10 +389,10 @@ async def _execute_tool_call(self, tool_call: FunctionCall, cancellation_token: if tool is None: raise ValueError(f"The tool '{tool_call.name}' is not available.") arguments = json.loads(tool_call.arguments) - result = await tool.run_json(arguments, cancellation_token) + result = await tool.run_json(arguments, cancellation_token, call_id=tool_call.id) return tool.return_value_as_string(result) - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: """Handle incoming messages and return a response.""" async for message in self.on_messages_stream(messages, cancellation_token): @@ -399,25 +401,22 @@ async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: raise AssertionError("The stream should have returned the final result.") async def on_messages_stream( - self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken - ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]: + self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken + ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]: """Handle incoming messages and return a response.""" await self._ensure_initialized() # Process all messages in sequence for message in messages: - if isinstance(message, (TextMessage, MultiModalMessage)): - await self.handle_text_message(str(message.content), cancellation_token) - elif isinstance(message, (StopMessage, HandoffMessage)): - await self.handle_text_message(message.content, cancellation_token) + await self.handle_incoming_message(message, cancellation_token) # Inner messages for tool calls - inner_messages: List[AgentEvent | ChatMessage] = [] + inner_messages: List[BaseAgentEvent | BaseChatMessage] = [] # Create and start a run run: Run = await cancellation_token.link_future( asyncio.ensure_future( - self._client.beta.threads.runs.create( + self._client.beta.threads.runs.create( # type: ignore[reportDeprecated] thread_id=self._thread_id, assistant_id=self._get_assistant_id, ) @@ -428,7 +427,7 @@ async def on_messages_stream( while True: run = await cancellation_token.link_future( asyncio.ensure_future( - self._client.beta.threads.runs.retrieve( + self._client.beta.threads.runs.retrieve( # type: ignore[reportDeprecated] thread_id=self._thread_id, run_id=run.id, ) @@ -481,7 +480,7 @@ async def on_messages_stream( # Submit tool outputs back to the run run = await cancellation_token.link_future( asyncio.ensure_future( - self._client.beta.threads.runs.submit_tool_outputs( + self._client.beta.threads.runs.submit_tool_outputs( # type: ignore[reportDeprecated] thread_id=self._thread_id, run_id=run.id, tool_outputs=[{"tool_call_id": t.call_id, "output": t.content} for t in tool_outputs], @@ -498,7 +497,7 @@ async def on_messages_stream( # Get messages after run completion assistant_messages: AsyncCursorPage[Message] = await cancellation_token.link_future( asyncio.ensure_future( - self._client.beta.threads.messages.list(thread_id=self._thread_id, order="desc", limit=1) + self._client.beta.threads.messages.list(thread_id=self._thread_id, order="desc", limit=1) # type: ignore[reportDeprecated] ) ) @@ -519,11 +518,24 @@ async def on_messages_stream( chat_message = TextMessage(source=self.name, content=text_content[0].text.value) yield Response(chat_message=chat_message, inner_messages=inner_messages) - async def handle_text_message(self, content: str, cancellation_token: CancellationToken) -> None: + async def handle_incoming_message(self, message: BaseChatMessage, cancellation_token: CancellationToken) -> None: """Handle regular text messages by adding them to the thread.""" + content: str | List[MessageContentPartParam] | None = None + llm_message = message.to_model_message() + if isinstance(llm_message.content, str): + content = llm_message.content + else: + content = [] + for c in llm_message.content: + if isinstance(c, str): + content.append(TextContentBlockParam(text=c, type="text")) + elif isinstance(c, Image): + content.append(ImageURLContentBlockParam(image_url=ImageURLParam(url=c.data_uri), type="image_url")) + else: + raise ValueError(f"Unsupported content type: {type(c)} in {message}") await cancellation_token.link_future( asyncio.ensure_future( - self._client.beta.threads.messages.create( + self._client.beta.threads.messages.create( # type: ignore[reportDeprecated] thread_id=self._thread_id, content=content, role="user", @@ -541,7 +553,7 @@ async def on_reset(self, cancellation_token: CancellationToken) -> None: while True: msgs: AsyncCursorPage[Message] = await cancellation_token.link_future( asyncio.ensure_future( - self._client.beta.threads.messages.list(self._thread_id, after=after, order="asc", limit=100) + self._client.beta.threads.messages.list(self._thread_id, after=after, order="asc", limit=100) # type: ignore[reportDeprecated] ) ) for msg in msgs.data: @@ -555,7 +567,7 @@ async def on_reset(self, cancellation_token: CancellationToken) -> None: for msg_id in new_message_ids: status: MessageDeleted = await cancellation_token.link_future( asyncio.ensure_future( - self._client.beta.threads.messages.delete(message_id=msg_id, thread_id=self._thread_id) + self._client.beta.threads.messages.delete(message_id=msg_id, thread_id=self._thread_id) # type: ignore[reportDeprecated] ) ) assert status.deleted is True @@ -591,7 +603,7 @@ async def on_upload_for_code_interpreter( # Update thread with the new files thread = await cancellation_token.link_future( - asyncio.ensure_future(self._client.beta.threads.retrieve(thread_id=self._thread_id)) + asyncio.ensure_future(self._client.beta.threads.retrieve(thread_id=self._thread_id)) # type: ignore[reportDeprecated] ) tool_resources: ToolResources = thread.tool_resources or ToolResources() code_interpreter: ToolResourcesCodeInterpreter = ( @@ -603,7 +615,7 @@ async def on_upload_for_code_interpreter( await cancellation_token.link_future( asyncio.ensure_future( - self._client.beta.threads.update( + self._client.beta.threads.update( # type: ignore[reportDeprecated] thread_id=self._thread_id, tool_resources=cast(thread_update_params.ToolResources, tool_resources.model_dump()), ) @@ -666,7 +678,7 @@ async def delete_assistant(self, cancellation_token: CancellationToken) -> None: if self._assistant is not None and not self._assistant_id: try: await cancellation_token.link_future( - asyncio.ensure_future(self._client.beta.assistants.delete(assistant_id=self._get_assistant_id)) + asyncio.ensure_future(self._client.beta.assistants.delete(assistant_id=self._get_assistant_id)) # type: ignore[reportDeprecated] ) self._assistant = None except Exception as e: diff --git a/python/packages/autogen-ext/src/autogen_ext/agents/web_surfer/_multimodal_web_surfer.py b/python/packages/autogen-ext/src/autogen_ext/agents/web_surfer/_multimodal_web_surfer.py index f4fb3abd10ea..e833a27ce3a4 100644 --- a/python/packages/autogen-ext/src/autogen_ext/agents/web_surfer/_multimodal_web_surfer.py +++ b/python/packages/autogen-ext/src/autogen_ext/agents/web_surfer/_multimodal_web_surfer.py @@ -24,7 +24,7 @@ import PIL.Image from autogen_agentchat.agents import BaseChatAgent from autogen_agentchat.base import Response -from autogen_agentchat.messages import AgentEvent, ChatMessage, MultiModalMessage, TextMessage +from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage, MultiModalMessage, TextMessage from autogen_agentchat.utils import content_to_str, remove_images from autogen_core import EVENT_LOGGER_NAME, CancellationToken, Component, ComponentModel, FunctionCall from autogen_core import Image as AGImage @@ -385,7 +385,7 @@ async def _set_debug_dir(self, debug_dir: str | None) -> None: ) @property - def produced_message_types(self) -> Sequence[type[ChatMessage]]: + def produced_message_types(self) -> Sequence[type[BaseChatMessage]]: return (MultiModalMessage,) async def on_reset(self, cancellation_token: CancellationToken) -> None: @@ -422,21 +422,19 @@ async def on_reset(self, cancellation_token: CancellationToken) -> None: ) ) - async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response: + async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response: async for message in self.on_messages_stream(messages, cancellation_token): if isinstance(message, Response): return message raise AssertionError("The stream should have returned the final result.") async def on_messages_stream( - self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken - ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]: + self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken + ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]: for chat_message in messages: - if isinstance(chat_message, TextMessage | MultiModalMessage): - self._chat_history.append(UserMessage(content=chat_message.content, source=chat_message.source)) - else: - raise ValueError(f"Unexpected message in MultiModalWebSurfer: {chat_message}") - self.inner_messages: List[AgentEvent | ChatMessage] = [] + self._chat_history.append(chat_message.to_model_message()) + + self.inner_messages: List[BaseAgentEvent | BaseChatMessage] = [] self.model_usage: List[RequestUsage] = [] try: content = await self._generate_reply(cancellation_token=cancellation_token) diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/_common.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/_common.py index 4b1259ef04ee..1ab1aa854b55 100644 --- a/python/packages/autogen-ext/src/autogen_ext/code_executors/_common.py +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/_common.py @@ -1,5 +1,6 @@ import inspect import re +import shutil from dataclasses import dataclass from pathlib import Path from textwrap import dedent, indent @@ -159,7 +160,13 @@ def lang_to_cmd(lang: str) -> str: if lang in ["shell"]: return "sh" if lang in ["pwsh", "powershell", "ps1"]: - return "pwsh" + # Check if pwsh is available, otherwise fall back to powershell + if shutil.which("pwsh") is not None: + return "pwsh" + elif shutil.which("powershell") is not None: + return "powershell" + else: + raise ValueError("Powershell or pwsh is not installed. Please install one of them.") else: raise ValueError(f"Unsupported language: {lang}") diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/azure/_azure_container_code_executor.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/azure/_azure_container_code_executor.py index 2a6f6c8ee4a7..17c4b16a2c15 100644 --- a/python/packages/autogen-ext/src/autogen_ext/code_executors/azure/_azure_container_code_executor.py +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/azure/_azure_container_code_executor.py @@ -4,6 +4,8 @@ import asyncio import os +import tempfile +import warnings from pathlib import Path from string import Template from typing import TYPE_CHECKING, Any, Callable, ClassVar, List, Optional, Protocol, Sequence, Union @@ -68,8 +70,13 @@ class ACADynamicSessionsCodeExecutor(CodeExecutor): timeout (int): The timeout for the execution of any single code block. Default is 60. work_dir (str): The working directory for the code execution. If None, a default working directory will be used. The default working - directory is the current directory ".". + directory is a temporal directory. functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]): A list of functions that are available to the code executor. Default is an empty list. + suppress_result_output bool: By default the executor will attach any result info in the execution response to the result outpu. Set this to True to prevent this. + session_id (str): The session id for the code execution (passed to Dynamic Sessions). If None, a new session id will be generated. Default is None. Note this value will be reset when calling `restart` + + .. note:: + Using the current directory (".") as working directory is deprecated. Using it will raise a deprecation warning. """ SUPPORTED_LANGUAGES: ClassVar[List[str]] = [ @@ -86,7 +93,7 @@ def __init__( pool_management_endpoint: str, credential: TokenProvider, timeout: int = 60, - work_dir: Union[Path, str] = Path("."), + work_dir: Union[Path, str, None] = None, functions: Sequence[ Union[ FunctionWithRequirements[Any, A], @@ -95,34 +102,48 @@ def __init__( ] ] = [], functions_module: str = "functions", + suppress_result_output: bool = False, + session_id: Optional[str] = None, ): if timeout < 1: raise ValueError("Timeout must be greater than or equal to 1.") - if isinstance(work_dir, str): - work_dir = Path(work_dir) + self._work_dir: Optional[Path] = None + self._temp_dir: Optional[tempfile.TemporaryDirectory[str]] = None + + # If a user specifies a working directory, use that + if work_dir is not None: + if isinstance(work_dir, str): + self._work_dir = Path(work_dir) + else: + self._work_dir = work_dir + # Create the directory if it doesn't exist + self._work_dir.mkdir(exist_ok=True, parents=True) + # If a user does not specify a working directory, use the default directory (tempfile.TemporaryDirectory) + else: + self._temp_dir = tempfile.TemporaryDirectory() + temp_dir_path = Path(self._temp_dir.name) + temp_dir_path.mkdir(exist_ok=True, parents=True) - if not functions_module.isidentifier(): - raise ValueError("Module name must be a valid Python identifier") + self._started = False + # Rest of initialization remains the same self._functions_module = functions_module - - work_dir.mkdir(exist_ok=True) - self._work_dir: Path = work_dir - self._timeout = timeout - self._functions = functions - self._func_code: str | None = None + self._func_code: Optional[str] = None + # Setup could take some time so we intentionally wait for the first code block to do it. if len(functions) > 0: self._setup_functions_complete = False else: self._setup_functions_complete = True + self._suppress_result_output = suppress_result_output + self._pool_management_endpoint = pool_management_endpoint self._access_token: str | None = None - self._session_id: str = str(uuid4()) + self._session_id: str = session_id or str(uuid4()) self._available_packages: set[str] | None = None self._credential: TokenProvider = credential # cwd needs to be set to /mnt/data to properly read uploaded files and download written files @@ -168,8 +189,21 @@ def timeout(self) -> int: @property def work_dir(self) -> Path: - """(Experimental) The working directory for the code execution.""" - return self._work_dir + # If a user specifies a working directory, use that + if self._work_dir is not None: + # If a user specifies the current directory, warn them that this is deprecated + if self._work_dir == Path("."): + warnings.warn( + "Using the current directory as work_dir is deprecated", + DeprecationWarning, + stacklevel=2, + ) + return self._work_dir + # If a user does not specify a working directory, use the default directory (tempfile.TemporaryDirectory) + elif self._temp_dir is not None: + return Path(self._temp_dir.name) + else: + raise RuntimeError("Working directory not properly initialized") def _construct_url(self, path: str) -> str: endpoint = self._pool_management_endpoint @@ -206,11 +240,18 @@ async def _setup_functions(self, cancellation_token: CancellationToken) -> None: flattened_packages = [item for sublist in lists_of_packages for item in sublist] required_packages = set(flattened_packages) + + if self._available_packages is None: + await self._populate_available_packages(cancellation_token) + if self._available_packages is not None: missing_pkgs = set(required_packages - self._available_packages) if len(missing_pkgs) > 0: raise ValueError(f"Packages unavailable in environment: {missing_pkgs}") + func_file = self.work_dir / f"{self._functions_module}.py" + func_file.write_text(self._func_code) + # Attempt to load the function file to check for syntax errors, imports etc. exec_result = await self._execute_code_dont_check_setup( [CodeBlock(code=self._func_code, language="python")], cancellation_token @@ -274,8 +315,8 @@ async def upload_files(self, files: List[Union[Path, str]], cancellation_token: timeout = aiohttp.ClientTimeout(total=float(self._timeout)) async with aiohttp.ClientSession(timeout=timeout) as client: for file in files: - file_path = os.path.join(self._work_dir, file) - if not os.path.isfile(file_path): + file_path = self.work_dir / file + if not file_path.is_file(): # TODO: what to do here? raise FileNotFoundError(f"{file} does not exist") @@ -335,8 +376,8 @@ async def download_files(self, files: List[Union[Path, str]], cancellation_token try: resp = await task resp.raise_for_status() - local_path = os.path.join(self._work_dir, file) - local_paths.append(local_path) + local_path = self.work_dir / file + local_paths.append(str(local_path)) async with await open_file(local_path, "wb") as f: await f.write(await resp.read()) except asyncio.TimeoutError as e: @@ -433,7 +474,8 @@ async def _execute_code_dont_check_setup( data = data["properties"] logs_all += data.get("stderr", "") + data.get("stdout", "") if "Success" in data["status"]: - logs_all += str(data["result"]) + if not self._suppress_result_output: + logs_all += str(data["result"]) elif "Failure" in data["status"]: exitcode = 1 @@ -452,7 +494,11 @@ async def _execute_code_dont_check_setup( return CodeResult(exit_code=exitcode, output=logs_all) async def restart(self) -> None: - """(Experimental) Restart the code executor.""" + """(Experimental) Restart the code executor. + + Resets the internal state of the executor by generating a new session ID and resetting the setup variables. + This causes the next code execution to reinitialize the environment and re-run any setup code. + """ self._session_id = str(uuid4()) self._setup_functions_complete = False self._access_token = None @@ -460,11 +506,17 @@ async def restart(self) -> None: self._setup_cwd_complete = False async def start(self) -> None: - """(Experimental) Start the code executor.""" + """(Experimental) Start the code executor. + + Marks the code executor as started.""" # No setup needed for this executor - pass + self._started = True async def stop(self) -> None: - """(Experimental) Stop the code executor.""" - # No cleanup needed for this executor - pass + """(Experimental) Stop the code executor. + + Stops the code executor after cleaning up the temporary working directory (if it was created).""" + if self._temp_dir is not None: + self._temp_dir.cleanup() + self._temp_dir = None + self._started = False diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py index 266effcd6ed2..cccf78cb6f36 100644 --- a/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py @@ -7,8 +7,11 @@ import logging import shlex import sys +import tempfile import uuid +import warnings from collections.abc import Sequence +from concurrent.futures import Future as ConcurrentFuture from hashlib import sha256 from pathlib import Path from typing import Any, Callable, ClassVar, Dict, List, Optional, ParamSpec, Tuple, Union @@ -23,6 +26,8 @@ from pydantic import BaseModel from typing_extensions import Self +from docker.types import DeviceRequest + from .._common import ( CommandLineCodeResult, build_python_functions_file, @@ -68,14 +73,15 @@ class DockerCommandLineCodeExecutorConfig(BaseModel): image: str = "python:3-slim" container_name: Optional[str] = None timeout: int = 60 - work_dir: str = "." # Stored as string, converted to Path - bind_dir: Optional[str] = None # Stored as string, converted to Path + work_dir: Optional[str] = None + bind_dir: Optional[str] = None auto_remove: bool = True stop_container: bool = True functions_module: str = "functions" extra_volumes: Dict[str, Dict[str, str]] = {} extra_hosts: Dict[str, str] = {} init_command: Optional[str] = None + delete_tmp_files: bool = False class DockerCommandLineCodeExecutor(CodeExecutor, Component[DockerCommandLineCodeExecutorConfig]): @@ -95,8 +101,7 @@ class DockerCommandLineCodeExecutor(CodeExecutor, Component[DockerCommandLineCod The executor executes the code blocks in the order they are received. Currently, the executor only supports Python and shell scripts. For Python code, use the language "python" for the code block. - For shell scripts, use the language "bash", "shell", or "sh" for the code - block. + For shell scripts, use the language "bash", "shell", "sh", "pwsh", "powershell", or "ps1" for the code block. Args: image (_type_, optional): Docker image to use for code execution. @@ -105,7 +110,7 @@ class DockerCommandLineCodeExecutor(CodeExecutor, Component[DockerCommandLineCod which is created. If None, will autogenerate a name. Defaults to None. timeout (int, optional): The timeout for code execution. Defaults to 60. work_dir (Union[Path, str], optional): The working directory for the code - execution. Defaults to Path("."). + execution. Defaults to temporary directory. bind_dir (Union[Path, str], optional): The directory that will be bound to the code executor container. Useful for cases where you want to spawn the container from within a container. Defaults to work_dir. @@ -114,6 +119,7 @@ class DockerCommandLineCodeExecutor(CodeExecutor, Component[DockerCommandLineCod stop_container (bool, optional): If true, will automatically stop the container when stop is called, when the context manager exits or when the Python process exits with atext. Defaults to True. + device_requests (Optional[List[DeviceRequest]], optional): A list of device request instances to add to the container for exposing GPUs (e.g., [docker.types.DeviceRequest(count=-1, capabilities=[['gpu']])]). Defaults to None. functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]): A list of functions that are available to the code executor. Default is an empty list. functions_module (str, optional): The name of the module that will be created to store the functions. Defaults to "functions". extra_volumes (Optional[Dict[str, Dict[str, str]]], optional): A dictionary of extra volumes (beyond the work_dir) to mount to the container; @@ -123,6 +129,11 @@ class DockerCommandLineCodeExecutor(CodeExecutor, Component[DockerCommandLineCod Example: extra_hosts = {"kubernetes.docker.internal": "host-gateway"} init_command (Optional[str], optional): A shell command to run before each shell operation execution. Defaults to None. Example: init_command="kubectl config use-context docker-hub" + delete_tmp_files (bool, optional): If true, will delete temporary files after execution. Defaults to False. + + .. note:: + Using the current directory (".") as working directory is deprecated. Using it will raise a deprecation warning. + """ component_config_schema = DockerCommandLineCodeExecutorConfig @@ -152,10 +163,11 @@ def __init__( container_name: Optional[str] = None, *, timeout: int = 60, - work_dir: Union[Path, str] = Path("."), + work_dir: Union[Path, str, None] = None, bind_dir: Optional[Union[Path, str]] = None, auto_remove: bool = True, stop_container: bool = True, + device_requests: Optional[List[DeviceRequest]] = None, functions: Sequence[ Union[ FunctionWithRequirements[Any, A], @@ -167,18 +179,27 @@ def __init__( extra_volumes: Optional[Dict[str, Dict[str, str]]] = None, extra_hosts: Optional[Dict[str, str]] = None, init_command: Optional[str] = None, + delete_tmp_files: bool = False, ): if timeout < 1: raise ValueError("Timeout must be greater than or equal to 1.") - if isinstance(work_dir, str): - work_dir = Path(work_dir) - work_dir.mkdir(exist_ok=True) - - if bind_dir is None: - bind_dir = work_dir - elif isinstance(bind_dir, str): - bind_dir = Path(bind_dir) + # Handle working directory logic + if work_dir is None: + self._work_dir = None + else: + if isinstance(work_dir, str): + work_dir = Path(work_dir) + # Emit a deprecation warning if the user is using the current directory as working directory + if work_dir.resolve() == Path.cwd().resolve(): + warnings.warn( + "Using the current directory as work_dir is deprecated.", + DeprecationWarning, + stacklevel=2, + ) + self._work_dir = work_dir + # Create the working directory if it doesn't exist + self._work_dir.mkdir(exist_ok=True, parents=True) if container_name is None: self.container_name = f"autogen-code-exec-{uuid.uuid4()}" @@ -186,8 +207,19 @@ def __init__( self.container_name = container_name self._timeout = timeout - self._work_dir: Path = work_dir - self._bind_dir: Path = bind_dir + + # Handle bind_dir + self._bind_dir: Optional[Path] = None + if bind_dir is not None: + self._bind_dir = Path(bind_dir) if isinstance(bind_dir, str) else bind_dir + else: + self._bind_dir = self._work_dir # Default to work_dir if not provided + + # Track temporary directory + self._temp_dir: Optional[tempfile.TemporaryDirectory[str]] = None + self._temp_dir_path: Optional[Path] = None + + self._started = False self._auto_remove = auto_remove self._stop_container = stop_container @@ -201,6 +233,8 @@ def __init__( self._extra_volumes = extra_volumes if extra_volumes is not None else {} self._extra_hosts = extra_hosts if extra_hosts is not None else {} self._init_command = init_command + self._delete_tmp_files = delete_tmp_files + self._device_requests = device_requests # Setup could take some time so we intentionally wait for the first code block to do it. if len(functions) > 0: @@ -210,26 +244,18 @@ def __init__( self._container: Container | None = None self._running = False - self._cancellation_tasks: List[asyncio.Task[None]] = [] + + self._loop: Optional[asyncio.AbstractEventLoop] = None + self._cancellation_futures: List[ConcurrentFuture[None]] = [] @property def timeout(self) -> int: """(Experimental) The timeout for code execution.""" return self._timeout - @property - def work_dir(self) -> Path: - """(Experimental) The working directory for the code execution.""" - return self._work_dir - - @property - def bind_dir(self) -> Path: - """(Experimental) The binding directory for the code execution container.""" - return self._bind_dir - async def _setup_functions(self, cancellation_token: CancellationToken) -> None: func_file_content = build_python_functions_file(self._functions) - func_file = self._work_dir / f"{self._functions_module}.py" + func_file = self.work_dir / f"{self._functions_module}.py" func_file.write_text(func_file_content) # Collect requirements @@ -282,7 +308,22 @@ async def _execute_command(self, command: List[str], cancellation_token: Cancell return output, exit_code except asyncio.CancelledError: # Schedule a task to kill the running command in the background. - self._cancellation_tasks.append(asyncio.create_task(self._kill_running_command(command))) + if self._loop and not self._loop.is_closed(): + try: + logging.debug(f"Scheduling kill command via run_coroutine_threadsafe on loop {self._loop!r}") + future: ConcurrentFuture[None] = asyncio.run_coroutine_threadsafe( + self._kill_running_command(command), self._loop + ) + self._cancellation_futures.append(future) + logging.debug(f"Kill command scheduled, future: {future!r}") + except RuntimeError as e: + logging.error(f"Failed to schedule kill command on loop {self._loop!r}: {e}") + except Exception as e: + logging.exception(f"Unexpected error scheduling kill command: {e}") + else: + logging.warning( + f"Cannot schedule kill command: Executor loop is not available or closed (loop: {self._loop!r})." + ) return "Code execution was cancelled.", 1 async def _execute_code_dont_check_setup( @@ -297,37 +338,72 @@ async def _execute_code_dont_check_setup( outputs: List[str] = [] files: List[Path] = [] last_exit_code = 0 - for code_block in code_blocks: - lang = code_block.language.lower() - code = silence_pip(code_block.code, lang) - - # Check if there is a filename comment - try: - filename = get_file_name_from_content(code, self._work_dir) - except ValueError: - outputs.append("Filename is not in the workspace") - last_exit_code = 1 - break - - if not filename: - filename = f"tmp_code_{sha256(code.encode()).hexdigest()}.{lang}" - - code_path = self._work_dir / filename - with code_path.open("w", encoding="utf-8") as fout: - fout.write(code) - files.append(code_path) - - command = ["timeout", str(self._timeout), lang_to_cmd(lang), filename] - - output, exit_code = await self._execute_command(command, cancellation_token) - outputs.append(output) - last_exit_code = exit_code - if exit_code != 0: - break + try: + for code_block in code_blocks: + lang = code_block.language.lower() + code = silence_pip(code_block.code, lang) + + # Check if there is a filename comment + try: + filename = get_file_name_from_content(code, self.work_dir) + except ValueError: + outputs.append("Filename is not in the workspace") + last_exit_code = 1 + break + + if not filename: + filename = f"tmp_code_{sha256(code.encode()).hexdigest()}.{lang}" + + code_path = self.work_dir / filename + with code_path.open("w", encoding="utf-8") as fout: + fout.write(code) + files.append(code_path) + + command = ["timeout", str(self._timeout), lang_to_cmd(lang), filename] + + output, exit_code = await self._execute_command(command, cancellation_token) + outputs.append(output) + last_exit_code = exit_code + if exit_code != 0: + break + finally: + if self._delete_tmp_files: + for file in files: + try: + file.unlink() + except (OSError, FileNotFoundError): + pass code_file = str(files[0]) if files else None return CommandLineCodeResult(exit_code=last_exit_code, output="".join(outputs), code_file=code_file) + @property + def work_dir(self) -> Path: + # If a user specifies a working directory, use that + if self._work_dir is not None: + # If a user specifies the current directory, warn them that this is deprecated + if self._work_dir == Path("."): + warnings.warn( + "Using the current directory as work_dir is deprecated.", + DeprecationWarning, + stacklevel=2, + ) + return self._work_dir + # If a user does not specify a working directory, use the default directory (tempfile.TemporaryDirectory) + elif self._temp_dir is not None: + return Path(self._temp_dir.name) + else: + raise RuntimeError("Working directory not properly initialized") + + @property + def bind_dir(self) -> Path: + # If the user specified a bind directory, return it + if self._bind_dir is not None: + return self._bind_dir + # Otherwise bind_dir is set to the current work_dir as default + else: + return self.work_dir + async def execute_code_blocks( self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken ) -> CommandLineCodeResult: @@ -345,10 +421,10 @@ async def execute_code_blocks( return await self._execute_code_dont_check_setup(code_blocks, cancellation_token) async def restart(self) -> None: + """(Experimental) Restart the Docker container code executor.""" if self._container is None or not self._running: raise ValueError("Container is not running. Must first be started with either start or a context manager.") - """(Experimental) Restart the code executor.""" await asyncio.to_thread(self._container.restart) # type: ignore if self._container.status != "running": self._running = False @@ -356,24 +432,80 @@ async def restart(self) -> None: raise ValueError(f"Failed to restart container. Logs: {logs_str}") async def stop(self) -> None: - """(Experimental) Stop the code executor.""" + """(Experimental) Stop the code executor. + Stops the Docker container and cleans up any temporary files (if they were created), along with the temporary directory. + The method first waits for all cancellation tasks to finish before stopping the container. Finally it marks the executor as not running. + If the container is not running, the method does nothing. + """ if not self._running: return + if self._temp_dir is not None: + self._temp_dir.cleanup() + self._temp_dir = None + client = docker.from_env() try: - container = await asyncio.to_thread(client.containers.get, self.container_name) - # Wait for all cancellation tasks to finish before stopping the container. - await asyncio.gather(*self._cancellation_tasks) - # Stop the container. + try: + container = await asyncio.to_thread(client.containers.get, self.container_name) + except NotFound: + logging.debug(f"Container {self.container_name} not found during stop...") + self._running = False + self._cancellation_futures.clear() + return + + if self._cancellation_futures: + if not self._loop or self._loop.is_closed(): + logging.warning( + f"Executor loop ({self._loop!r}) is closed or unavailable. Cannot reliably wait for " + f"{len(self._cancellation_futures)} cancellation futures." + ) + self._cancellation_futures.clear() + else: + # concurrent.futures.Future -> asyncio.Future + asyncio_futures = [asyncio.wrap_future(f, loop=self._loop) for f in self._cancellation_futures] + + if asyncio_futures: + logging.debug( + f"Waiting for {len(asyncio_futures)} cancellation futures to complete on loop {self._loop!r}..." + ) + results = await asyncio.gather(*asyncio_futures, return_exceptions=True) + for i, result in enumerate(results): + original_future = self._cancellation_futures[i] + if isinstance(result, Exception): + logging.warning(f"Cancellation future {original_future!r} failed: {result}") + else: + logging.debug(f"Cancellation future {original_future!r} completed successfully.") + else: + logging.debug("No valid cancellation futures to await.") + + self._cancellation_futures.clear() + + logging.debug(f"Stopping container {self.container_name}...") await asyncio.to_thread(container.stop) - except NotFound: - pass + logging.debug(f"Container {self.container_name} stopped.") + + except DockerException as e: + logging.error(f"Docker error while stopping container {self.container_name}: {e}") + except Exception as e: + logging.exception(f"Unexpected error during stop operation for container {self.container_name}: {e}") finally: self._running = False + self._cancellation_futures.clear() async def start(self) -> None: + """(Experimental) Start the code executor. + + This method sets the working environment variables, connects to Docker and starts the code executor. + If no working directory was provided to the code executor, it creates a temporary directory and sets it as the code executor working directory. + """ + + if self._work_dir is None and self._temp_dir is None: + self._temp_dir = tempfile.TemporaryDirectory() + self._temp_dir_path = Path(self._temp_dir.name) + self._temp_dir_path.mkdir(exist_ok=True) + # Start a container from the image, read to exec commands later try: client = docker.from_env() @@ -397,6 +529,13 @@ async def start(self) -> None: shell_command = "/bin/sh" command = ["-c", f"{(self._init_command)};exec {shell_command}"] if self._init_command else None + # Check if a container with the same name already exists and remove it + try: + existing_container = await asyncio.to_thread(client.containers.get, self.container_name) + await asyncio.to_thread(existing_container.remove, force=True) + except NotFound: + pass + self._container = await asyncio.to_thread( client.containers.create, self._image, @@ -406,9 +545,10 @@ async def start(self) -> None: tty=True, detach=True, auto_remove=self._auto_remove, - volumes={str(self._bind_dir.resolve()): {"bind": "/workspace", "mode": "rw"}, **self._extra_volumes}, + volumes={str(self.bind_dir.resolve()): {"bind": "/workspace", "mode": "rw"}, **self._extra_volumes}, working_dir="/workspace", extra_hosts=self._extra_hosts, + device_requests=self._device_requests, ) await asyncio.to_thread(self._container.start) @@ -426,6 +566,10 @@ async def cleanup() -> None: logs_str = self._container.logs().decode("utf-8") raise ValueError(f"Failed to start container from image {self._image}. Logs: {logs_str}") + self._loop = asyncio.get_running_loop() + self._cancellation_futures = [] + logging.debug(f"Executor started, associated with event loop: {self._loop!r}") + self._running = True def _to_config(self) -> DockerCommandLineCodeExecutorConfig: @@ -437,7 +581,7 @@ def _to_config(self) -> DockerCommandLineCodeExecutorConfig: image=self._image, container_name=self.container_name, timeout=self._timeout, - work_dir=str(self._work_dir), + work_dir=str(self._work_dir) if self._work_dir else None, bind_dir=str(self._bind_dir) if self._bind_dir else None, auto_remove=self._auto_remove, stop_container=self._stop_container, @@ -445,18 +589,19 @@ def _to_config(self) -> DockerCommandLineCodeExecutorConfig: extra_volumes=self._extra_volumes, extra_hosts=self._extra_hosts, init_command=self._init_command, + delete_tmp_files=self._delete_tmp_files, ) @classmethod def _from_config(cls, config: DockerCommandLineCodeExecutorConfig) -> Self: """(Experimental) Create a component from a config object.""" - bind_dir = Path(config.bind_dir) if config.bind_dir else None + return cls( image=config.image, container_name=config.container_name, timeout=config.timeout, - work_dir=Path(config.work_dir), - bind_dir=bind_dir, + work_dir=Path(config.work_dir) if config.work_dir else None, + bind_dir=Path(config.bind_dir) if config.bind_dir else None, auto_remove=config.auto_remove, stop_container=config.stop_container, functions=[], # Functions not restored from config @@ -464,4 +609,5 @@ def _from_config(cls, config: DockerCommandLineCodeExecutorConfig) -> Self: extra_volumes=config.extra_volumes, extra_hosts=config.extra_hosts, init_command=config.init_command, + delete_tmp_files=config.delete_tmp_files, ) diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/docker_jupyter/__init__.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker_jupyter/__init__.py new file mode 100644 index 000000000000..549c178f16b1 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker_jupyter/__init__.py @@ -0,0 +1,10 @@ +from ._docker_jupyter import DockerJupyterCodeExecutor, DockerJupyterCodeResult +from ._jupyter_server import DockerJupyterServer, JupyterClient, JupyterKernelClient + +__all__ = [ + "DockerJupyterCodeExecutor", + "DockerJupyterServer", + "JupyterClient", + "JupyterKernelClient", + "DockerJupyterCodeResult", +] diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/docker_jupyter/_docker_jupyter.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker_jupyter/_docker_jupyter.py new file mode 100644 index 000000000000..a8f370d159ee --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker_jupyter/_docker_jupyter.py @@ -0,0 +1,300 @@ +import asyncio +import base64 +import json +import os +import tempfile +import uuid +from dataclasses import dataclass +from pathlib import Path +from types import TracebackType +from typing import List, Optional, Union + +from autogen_core import CancellationToken, Component +from autogen_core.code_executor import CodeBlock, CodeExecutor, CodeResult +from autogen_ext.code_executors._common import silence_pip +from pydantic import BaseModel +from typing_extensions import Self + +from ._jupyter_server import JupyterClient, JupyterConnectable, JupyterConnectionInfo, JupyterKernelClient + + +@dataclass +class DockerJupyterCodeResult(CodeResult): + """(Experimental) A code result class for IPython code executor.""" + + output_files: list[Path] + + +class DockerJupyterCodeExecutorConfig(BaseModel): + """Configuration for JupyterCodeExecutor""" + + jupyter_server: Union[JupyterConnectable, JupyterConnectionInfo] + kernel_name: str = "python3" + timeout: int = 60 + output_dir: Optional[Union[Path, str]] = None + + class Config: + arbitrary_types_allowed = True + + +class DockerJupyterCodeExecutor(CodeExecutor, Component[DockerJupyterCodeExecutorConfig]): + """(Experimental) A code executor class that executes code statefully using + a Jupyter server supplied to this class. + + Each execution is stateful and can access variables created from previous + executions in the same session. + + To use this, you need to install the following dependencies: + + .. code-block:: shell + + pip install "autogen-ext[docker-jupyter-executor]" + + Args: + jupyter_server (Union[JupyterConnectable, JupyterConnectionInfo]): The Jupyter server to use. + kernel_name (str): The kernel name to use. Make sure it is installed. + By default, it is "python3". + timeout (int): The timeout for code execution, by default 60. + output_dir (str): The directory to save output files, by default None. + + Example of using it directly: + + .. code-block:: python + + import asyncio + from autogen_core import CancellationToken + from autogen_core.code_executor import CodeBlock + from autogen_ext.code_executors.docker_jupyter import DockerJupyterCodeExecutor, DockerJupyterServer + + + async def main() -> None: + async with DockerJupyterServer() as jupyter_server: + async with DockerJupyterCodeExecutor(jupyter_server=jupyter_server) as executor: + code_blocks = [CodeBlock(code="print('hello world!')", language="python")] + code_result = await executor.execute_code_blocks(code_blocks, cancellation_token=CancellationToken()) + print(code_result) + + + asyncio.run(main()) + + Example of using it with your own jupyter image: + + .. code-block:: python + + import asyncio + from autogen_core import CancellationToken + from autogen_core.code_executor import CodeBlock + from autogen_ext.code_executors.docker_jupyter import DockerJupyterCodeExecutor, DockerJupyterServer + + + async def main() -> None: + async with DockerJupyterServer(custom_image_name="your_custom_images_name", expose_port=8888) as jupyter_server: + async with DockerJupyterCodeExecutor(jupyter_server=jupyter_server) as executor: + code_blocks = [CodeBlock(code="print('hello world!')", language="python")] + code_result = await executor.execute_code_blocks(code_blocks, cancellation_token=CancellationToken()) + print(code_result) + + + asyncio.run(main()) + + Example of using it with :class:`~autogen_ext.tools.code_execution.PythonCodeExecutionTool`: + + .. code-block:: python + + import asyncio + from autogen_agentchat.agents import AssistantAgent + from autogen_ext.code_executors.docker_jupyter import DockerJupyterCodeExecutor, DockerJupyterServer + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_ext.tools.code_execution import PythonCodeExecutionTool + + + async def main() -> None: + async with DockerJupyterServer() as jupyter_server: + async with DockerJupyterCodeExecutor(jupyter_server=jupyter_server) as executor: + tool = PythonCodeExecutionTool(executor) + model_client = OpenAIChatCompletionClient(model="gpt-4o") + agent = AssistantAgent("assistant", model_client=model_client, tools=[tool]) + result = await agent.run(task="What is the 10th Fibonacci number? Use Python to calculate it.") + print(result) + + + asyncio.run(main()) + + Example of using it inside a :class:`~autogen_agentchat.agents._code_executor_agent.CodeExecutorAgent`: + + .. code-block:: python + + import asyncio + from autogen_agentchat.agents import CodeExecutorAgent + from autogen_agentchat.messages import TextMessage + from autogen_ext.code_executors.docker_jupyter import DockerJupyterCodeExecutor, DockerJupyterServer + from autogen_core import CancellationToken + + + async def main() -> None: + async with DockerJupyterServer() as jupyter_server: + async with DockerJupyterCodeExecutor(jupyter_server=jupyter_server) as executor: + code_executor_agent = CodeExecutorAgent("code_executor", code_executor=executor) + task = TextMessage( + content='''Here is some code + ```python + print('Hello world') + ``` + ''', + source="user", + ) + response = await code_executor_agent.on_messages([task], CancellationToken()) + print(response.chat_message) + + + asyncio.run(main()) + + """ + + component_config_schema = DockerJupyterCodeExecutorConfig + component_provider_override = "autogen_ext.code_executors.docker_jupyter.DockerJupyterCodeExecutor" + + def __init__( + self, + jupyter_server: Union[JupyterConnectable, JupyterConnectionInfo], + kernel_name: str = "python3", + timeout: int = 60, + output_dir: Path | None = None, + ): + if timeout < 1: + raise ValueError("Timeout must be greater than or equal to 1.") + + if isinstance(jupyter_server, JupyterConnectable): + self._connection_info = jupyter_server.connection_info + elif isinstance(jupyter_server, JupyterConnectionInfo): + self._connection_info = jupyter_server + else: + raise ValueError("jupyter_server must be a JupyterConnectable or JupyterConnectionInfo.") + + self._output_dir = output_dir or getattr(jupyter_server, "_bind_dir", None) + if not self._output_dir: + with tempfile.TemporaryDirectory() as temp_dir: + self._output_dir = Path(temp_dir) + self._output_dir.mkdir(exist_ok=True) + + self._jupyter_client = JupyterClient(self._connection_info) + + self._kernel_name = kernel_name + self._timeout = timeout + self._async_jupyter_kernel_client: Optional[JupyterKernelClient] = None + self._kernel_id: Optional[str] = None + + async def _ensure_async_kernel_client(self) -> JupyterKernelClient: + """Ensure that an async kernel client exists and return it.""" + if self._kernel_id is None: + await self.start() + assert self._kernel_id is not None + if self._async_jupyter_kernel_client is None: + self._async_jupyter_kernel_client = await self._jupyter_client.get_kernel_client(self._kernel_id) + return self._async_jupyter_kernel_client + + async def execute_code_blocks( + self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken + ) -> DockerJupyterCodeResult: + """(Experimental) Execute a list of code blocks and return the result. + + This method executes a list of code blocks as cells in the Jupyter kernel. + See: https://jupyter-client.readthedocs.io/en/stable/messaging.html + for the message protocol. + + Args: + code_blocks (List[CodeBlock]): A list of code blocks to execute. + + Returns: + DockerJupyterCodeResult: The result of the code execution. + """ + kernel_client = await self._ensure_async_kernel_client() + # Wait for kernel to be ready using async client + is_ready = await kernel_client.wait_for_ready(timeout_seconds=self._timeout) + if not is_ready: + return DockerJupyterCodeResult(exit_code=1, output="ERROR: Kernel not ready", output_files=[]) + + outputs: List[str] = [] + output_files: List[Path] = [] + for code_block in code_blocks: + code = silence_pip(code_block.code, code_block.language) + # Execute code using async client + exec_task = asyncio.create_task(kernel_client.execute(code, timeout_seconds=self._timeout)) + cancellation_token.link_future(exec_task) + result = await exec_task + if result.is_ok: + outputs.append(result.output) + for data in result.data_items: + if data.mime_type == "image/png": + path = self._save_image(data.data) + outputs.append(path) + output_files.append(Path(path)) + elif data.mime_type == "text/html": + path = self._save_html(data.data) + outputs.append(path) + output_files.append(Path(path)) + else: + outputs.append(json.dumps(data.data)) + else: + existing_output = "\n".join([str(output) for output in outputs]) + return DockerJupyterCodeResult( + exit_code=1, output=existing_output + "\nERROR: " + result.output, output_files=output_files + ) + return DockerJupyterCodeResult( + exit_code=0, output="\n".join([str(output) for output in outputs]), output_files=output_files + ) + + async def restart(self) -> None: + """(Experimental) Restart a new session.""" + # Use async client to restart kernel + if self._kernel_id is not None: + await self._jupyter_client.restart_kernel(self._kernel_id) + # Reset the clients to force recreation + if self._async_jupyter_kernel_client is not None: + await self._async_jupyter_kernel_client.stop() + self._async_jupyter_kernel_client = None + + async def start(self) -> None: + """(Experimental) Start a new session.""" + available_kernels = await self._jupyter_client.list_kernel_specs() + if self._kernel_name not in available_kernels["kernelspecs"]: + raise ValueError(f"Kernel {self._kernel_name} is not installed.") + self._kernel_id = await self._jupyter_client.start_kernel(self._kernel_name) + + def _save_image(self, image_data_base64: str) -> str: + """Save image data to a file.""" + image_data = base64.b64decode(image_data_base64) + filename = f"{uuid.uuid4().hex}.png" + path = os.path.join(str(self._output_dir), filename) + with open(path, "wb") as f: + f.write(image_data) + return os.path.abspath(path) + + def _save_html(self, html_data: str) -> str: + """Save html data to a file.""" + filename = f"{uuid.uuid4().hex}.html" + path = os.path.join(str(self._output_dir), filename) + with open(path, "w") as f: + f.write(html_data) + return os.path.abspath(path) + + async def stop(self) -> None: + """Stop the kernel.""" + if self._kernel_id is not None: + await self._jupyter_client.delete_kernel(self._kernel_id) + if self._async_jupyter_kernel_client is not None: + await self._async_jupyter_kernel_client.stop() + self._async_jupyter_kernel_client = None + await self._jupyter_client.close() + + async def __aenter__(self) -> Self: + await self.start() + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: + await self.stop() diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/docker_jupyter/_jupyter_server.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker_jupyter/_jupyter_server.py new file mode 100644 index 000000000000..be7f15e2c939 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker_jupyter/_jupyter_server.py @@ -0,0 +1,430 @@ +import asyncio +import atexit +import datetime +import io +import json +import logging +import os +import secrets +import uuid +from dataclasses import dataclass +from pathlib import Path +from time import sleep +from types import TracebackType +from typing import Any, Dict, List, Optional, Protocol, Type, Union, cast, runtime_checkable + +import aiohttp +import docker +import docker.errors +import requests +import websockets +from requests.adapters import HTTPAdapter, Retry +from typing_extensions import Self + + +@dataclass +class JupyterConnectionInfo: + """(Experimental)""" + + host: str + """`str` - Host of the Jupyter gateway server""" + use_https: bool + """`bool` - Whether to use HTTPS""" + port: Optional[int] = None + """`Optional[int]` - Port of the Jupyter gateway server. If None, the default port is used""" + token: Optional[str] = None + """`Optional[str]` - Token for authentication. If None, no token is used""" + + +@runtime_checkable +class JupyterConnectable(Protocol): + """(Experimental)""" + + @property + def connection_info(self) -> JupyterConnectionInfo: + """Return the connection information for this connectable.""" + ... + + +class JupyterClient: + def __init__(self, connection_info: JupyterConnectionInfo): + """(Experimental) A client for communicating with a Jupyter gateway server. + + Args: + connection_info (JupyterConnectionInfo): Connection information + """ + self._connection_info = connection_info + self._session = requests.Session() + retries = Retry(total=5, backoff_factor=0.1) + self._session.mount("http://", HTTPAdapter(max_retries=retries)) + # Create aiohttp session for async requests + self._async_session: aiohttp.ClientSession | None = None + + async def _ensure_async_session(self) -> aiohttp.ClientSession: + if self._async_session is None: + self._async_session = aiohttp.ClientSession() + return self._async_session + + def _get_headers(self) -> Dict[str, str]: + if self._connection_info.token is None: + return {} + return {"Authorization": f"token {self._connection_info.token}"} + + def _get_api_base_url(self) -> str: + protocol = "https" if self._connection_info.use_https else "http" + port = f":{self._connection_info.port}" if self._connection_info.port else "" + return f"{protocol}://{self._connection_info.host}{port}" + + def _get_ws_base_url(self) -> str: + port = f":{self._connection_info.port}" if self._connection_info.port else "" + return f"ws://{self._connection_info.host}{port}" + + async def list_kernel_specs(self) -> Dict[str, Dict[str, str]]: + response = self._session.get(f"{self._get_api_base_url()}/api/kernelspecs", headers=self._get_headers()) + return cast(Dict[str, Dict[str, str]], response.json()) + + async def list_kernels(self) -> List[Dict[str, str]]: + response = self._session.get(f"{self._get_api_base_url()}/api/kernels", headers=self._get_headers()) + return cast(List[Dict[str, str]], response.json()) + + async def start_kernel(self, kernel_spec_name: str) -> str: + """Start a new kernel asynchronously. + + Args: + kernel_spec_name (str): Name of the kernel spec to start + + Returns: + str: ID of the started kernel + """ + session = await self._ensure_async_session() + async with session.post( + f"{self._get_api_base_url()}/api/kernels", + headers=self._get_headers(), + json={"name": kernel_spec_name}, + ) as response: + data = await response.json() + return cast(str, data["id"]) + + async def delete_kernel(self, kernel_id: str) -> None: + session = await self._ensure_async_session() + async with session.delete( + f"{self._get_api_base_url()}/api/kernels/{kernel_id}", headers=self._get_headers() + ) as response: + response.raise_for_status() + + async def restart_kernel(self, kernel_id: str) -> None: + session = await self._ensure_async_session() + async with session.post( + f"{self._get_api_base_url()}/api/kernels/{kernel_id}/restart", headers=self._get_headers() + ) as response: + response.raise_for_status() + + async def get_kernel_client(self, kernel_id: str) -> "JupyterKernelClient": + ws_url = f"{self._get_ws_base_url()}/api/kernels/{kernel_id}/channels" + # Using websockets library for async websocket connections + ws = await websockets.connect(ws_url, additional_headers=self._get_headers()) + return JupyterKernelClient(ws) + + async def close(self) -> None: + """Close the async session""" + if self._async_session is not None: + await self._async_session.close() + self._async_session = None + self._session.close() + + +@dataclass +class DataItem: + mime_type: str + data: str + + +@dataclass +class ExecutionResult: + is_ok: bool + output: str + data_items: List[DataItem] + + +class JupyterKernelClient: + """An asynchronous client for communicating with a Jupyter kernel.""" + + def __init__(self, websocket: websockets.ClientConnection) -> None: + self._session_id = uuid.uuid4().hex + self._websocket = websocket + + async def __aenter__(self) -> Self: + return self + + async def __aexit__( + self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType] + ) -> None: + await self.stop() + + async def stop(self) -> None: + await self._websocket.close() + + async def _send_message(self, *, content: Dict[str, Any], channel: str, message_type: str) -> str: + timestamp = datetime.datetime.now().isoformat() + message_id = uuid.uuid4().hex + message = { + "header": { + "username": "autogen", + "version": "5.0", + "session": self._session_id, + "msg_id": message_id, + "msg_type": message_type, + "date": timestamp, + }, + "parent_header": {}, + "channel": channel, + "content": content, + "metadata": {}, + "buffers": {}, + } + await self._websocket.send(json.dumps(message)) + return message_id + + async def _receive_message(self, timeout_seconds: Optional[float]) -> Optional[Dict[str, Any]]: + try: + if timeout_seconds is not None: + data = await asyncio.wait_for(self._websocket.recv(), timeout=timeout_seconds) + else: + data = await self._websocket.recv() + if isinstance(data, bytes): + return cast(Dict[str, Any], json.loads(data.decode("utf-8"))) + return cast(Dict[str, Any], json.loads(data)) + except asyncio.TimeoutError: + return None + + async def wait_for_ready(self, timeout_seconds: Optional[float] = None) -> bool: + message_id = await self._send_message(content={}, channel="shell", message_type="kernel_info_request") + while True: + message = await self._receive_message(timeout_seconds) + # This means we timed out with no new messages. + if message is None: + return False + if ( + message.get("parent_header", {}).get("msg_id") == message_id + and message["msg_type"] == "kernel_info_reply" + ): + return True + + async def execute(self, code: str, timeout_seconds: Optional[float] = None) -> ExecutionResult: + message_id = await self._send_message( + content={ + "code": code, + "silent": False, + "store_history": True, + "user_expressions": {}, + "allow_stdin": False, + "stop_on_error": True, + }, + channel="shell", + message_type="execute_request", + ) + + text_output: List[str] = [] + data_output: List[DataItem] = [] + while True: + message = await self._receive_message(timeout_seconds) + if message is None: + return ExecutionResult( + is_ok=False, output="ERROR: Timeout waiting for output from code block.", data_items=[] + ) + + # Ignore messages that are not for this execution. + if message.get("parent_header", {}).get("msg_id") != message_id: + continue + + msg_type = message["msg_type"] + content = message["content"] + if msg_type in ["execute_result", "display_data"]: + for data_type, data in content["data"].items(): + if data_type == "text/plain": + text_output.append(data) + elif data_type.startswith("image/") or data_type == "text/html": + data_output.append(DataItem(mime_type=data_type, data=data)) + else: + text_output.append(json.dumps(data)) + elif msg_type == "stream": + text_output.append(content["text"]) + elif msg_type == "error": + # Output is an error. + return ExecutionResult( + is_ok=False, + output=f"ERROR: {content['ename']}: {content['evalue']}\n{content['traceback']}", + data_items=[], + ) + if msg_type == "status" and content["execution_state"] == "idle": + break + return ExecutionResult( + is_ok=True, output="\n".join([str(output) for output in text_output]), data_items=data_output + ) + + +class DockerJupyterServer(JupyterConnectable): + DEFAULT_DOCKERFILE = """FROM quay.io/jupyter/docker-stacks-foundation + + SHELL ["/bin/bash", "-o", "pipefail", "-c"] + + USER ${NB_UID} + RUN mamba install --yes jupyter_kernel_gateway ipykernel && \ + mamba clean --all -f -y && \ + fix-permissions "${CONDA_DIR}" && \ + fix-permissions "/home/${NB_USER}" + + ENV TOKEN="UNSET" + CMD python -m jupyter kernelgateway --KernelGatewayApp.ip=0.0.0.0 \ + --KernelGatewayApp.port=8888 \ + --KernelGatewayApp.auth_token="${TOKEN}" \ + --JupyterApp.answer_yes=true \ + --JupyterWebsocketPersonality.list_kernels=true + + EXPOSE 8888 + + WORKDIR "${HOME}" + """ + + class GenerateToken: + pass + + def __init__( + self, + *, + custom_image_name: Optional[str] = None, + container_name: Optional[str] = None, + auto_remove: bool = True, + stop_container: bool = True, + docker_env: Optional[Dict[str, str]] = None, + expose_port: int = 8888, + token: Optional[Union[str, GenerateToken]] = None, + work_dir: Union[Path, str] = "/workspace", + bind_dir: Optional[Union[Path, str]] = None, + ): + """Start a Jupyter kernel gateway server in a Docker container. + + Args: + custom_image_name: Custom Docker image to use. If None, builds and uses bundled image. + container_name: Name for the Docker container. Auto-generated if None. + auto_remove: If True, container will be deleted when stopped. + stop_container: If True, container stops on program exit or when context manager exits. + docker_env: Additional environment variables for the container. + expose_port: Port to expose for Jupyter connection. + token: Authentication token. If GenerateToken, creates random token. Empty for no auth. + work_dir: Working directory inside the container. + bind_dir: Local directory to bind to container's work_dir. + """ + # Generate container name if not provided + container_name = container_name or f"autogen-jupyterkernelgateway-{uuid.uuid4()}" + + # Initialize Docker client + client = docker.from_env() + # Set up bind directory if specified + self._bind_dir: Optional[Path] = None + if bind_dir: + self._bind_dir = Path(bind_dir) if isinstance(bind_dir, str) else bind_dir + self._bind_dir.mkdir(exist_ok=True) + os.chmod(bind_dir, 0o777) + + # Determine and prepare Docker image + image_name = custom_image_name or "autogen-jupyterkernelgateway" + if not custom_image_name: + try: + client.images.get(image_name) + except docker.errors.ImageNotFound: + # Build default image if not found + here = Path(__file__).parent + dockerfile = io.BytesIO(self.DEFAULT_DOCKERFILE.encode("utf-8")) + logging.info(f"Building image {image_name}...") + client.images.build(path=str(here), fileobj=dockerfile, tag=image_name) + logging.info(f"Image {image_name} built successfully") + else: + # Verify custom image exists + try: + client.images.get(image_name) + except docker.errors.ImageNotFound as err: + raise ValueError(f"Custom image {image_name} does not exist") from err + if docker_env is None: + docker_env = {} + if token is None: + token = DockerJupyterServer.GenerateToken() + # Set up authentication token + self._token = secrets.token_hex(32) if isinstance(token, DockerJupyterServer.GenerateToken) else token + + # Prepare environment variables + env = {"TOKEN": self._token} + env.update(docker_env) + + # Define volume configuration if bind directory is specified + volumes = {str(self._bind_dir): {"bind": str(work_dir), "mode": "rw"}} if self._bind_dir else None + + # Start the container + container = client.containers.run( + image_name, + detach=True, + auto_remove=auto_remove, + environment=env, + publish_all_ports=True, + name=container_name, + volumes=volumes, + working_dir=str(work_dir), + ) + + # Wait for container to be ready + self._wait_for_ready(container) + + # Store container information + self._container = container + self._port = int(container.ports[f"{expose_port}/tcp"][0]["HostPort"]) + self._container_id = container.id + self._expose_port = expose_port + + if self._container_id is None: + raise ValueError("Failed to obtain container id.") + + # Define cleanup function + def cleanup() -> None: + try: + assert self._container_id is not None + inner_container = client.containers.get(self._container_id) + inner_container.stop() + except docker.errors.NotFound: + pass + atexit.unregister(cleanup) + + # Register cleanup if container should be stopped automatically + if stop_container: + atexit.register(cleanup) + + self._cleanup_func = cleanup + self._stop_container = stop_container + + @property + def connection_info(self) -> JupyterConnectionInfo: + return JupyterConnectionInfo(host="127.0.0.1", use_https=False, port=self._port, token=self._token) + + def _wait_for_ready(self, container: Any, timeout: int = 60, stop_time: float = 0.1) -> None: + elapsed_time = 0.0 + while container.status != "running" and elapsed_time < timeout: + sleep(stop_time) + elapsed_time += stop_time + container.reload() + continue + if container.status != "running": + raise ValueError("Container failed to start") + + async def stop(self) -> None: + loop = asyncio.get_event_loop() + await loop.run_in_executor(None, self._cleanup_func) + + async def get_client(self) -> JupyterClient: + return JupyterClient(self.connection_info) + + async def __aenter__(self) -> Self: + return self + + async def __aexit__( + self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType] + ) -> None: + await self.stop() diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py index 385c26b6c9e3..2476b5a3349f 100644 --- a/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py @@ -3,7 +3,9 @@ import json import re import sys +import tempfile import uuid +import warnings from dataclasses import dataclass from pathlib import Path @@ -15,6 +17,9 @@ else: from typing_extensions import Self +from contextlib import AbstractAsyncContextManager +from typing import Optional, Union + from autogen_core import CancellationToken from autogen_core.code_executor import CodeBlock, CodeExecutor, CodeResult from nbclient import NotebookClient @@ -37,7 +42,7 @@ class JupyterCodeExecutorConfig(BaseModel): kernel_name: str = "python3" timeout: int = 60 - output_dir: str = "." + output_dir: Optional[str] = None class JupyterCodeExecutor(CodeExecutor, Component[JupyterCodeExecutorConfig]): @@ -121,7 +126,11 @@ async def main() -> None: Args: kernel_name (str): The kernel name to use. By default, "python3". timeout (int): The timeout for code execution, by default 60. - output_dir (Path): The directory to save output files, by default ".". + output_dir (Path): The directory to save output files, by default a temporary directory. + + + .. note:: + Using the current directory (".") as output directory is deprecated. Using it will raise a deprecation warning. """ component_config_schema = JupyterCodeExecutorConfig @@ -131,21 +140,24 @@ def __init__( self, kernel_name: str = "python3", timeout: int = 60, - output_dir: Path = Path("."), + output_dir: Optional[Union[Path, str]] = None, ): if timeout < 1: raise ValueError("Timeout must be greater than or equal to 1.") + self._output_dir: Path = Path(tempfile.mkdtemp()) if output_dir is None else Path(output_dir) + self._output_dir.mkdir(exist_ok=True, parents=True) + + self._temp_dir: Optional[tempfile.TemporaryDirectory[str]] = None + self._temp_dir_path: Optional[Path] = None + + self._started = False + self._kernel_name = kernel_name self._timeout = timeout - self._output_dir = output_dir - # TODO: Forward arguments perhaps? - self._client = NotebookClient( - nb=nbformat.new_notebook(), # type: ignore - kernel_name=self._kernel_name, - timeout=self._timeout, - allow_errors=True, - ) + + self._client: Optional[NotebookClient] = None + self.kernel_context: Optional[AbstractAsyncContextManager[None]] = None async def execute_code_blocks( self, code_blocks: list[CodeBlock], cancellation_token: CancellationToken @@ -230,6 +242,8 @@ async def _execute_code_block( async def _execute_cell(self, cell: NotebookNode) -> NotebookNode: # Temporary push cell to nb as async_execute_cell expects it. But then we want to remove it again as cells can take up significant amount of memory (especially with images) + if not self._client: + raise RuntimeError("Executor must be started before executing cells") self._client.nb.cells.append(cell) output = await self._client.async_execute_cell( cell, @@ -257,20 +271,65 @@ async def restart(self) -> None: await self.start() async def start(self) -> None: + """(Experimental) Start the code executor. + + Initializes the Jupyter Notebook execution environment by creating a new notebook and setting it up with the specified Jupyter Kernel. + Marks the executor as started, allowing for code execution. + This method should be called before executing any code blocks. + """ + if self._started: + return + + notebook: NotebookNode = nbformat.new_notebook() # type: ignore + + self._client = NotebookClient( + nb=notebook, + kernel_name=self._kernel_name, + timeout=self._timeout, + allow_errors=True, + ) + self.kernel_context = self._client.async_setup_kernel() await self.kernel_context.__aenter__() + self._started = True + async def stop(self) -> None: - """Stop the kernel.""" - await self.kernel_context.__aexit__(None, None, None) + """(Experimental) Stop the code executor. + + Terminates the Jupyter Notebook execution by exiting the kernel context and cleaning up the associated resources.""" + if not self._started: + return + + if self.kernel_context is not None: + await self.kernel_context.__aexit__(None, None, None) + self.kernel_context = None + + self._client = None + self._started = False def _to_config(self) -> JupyterCodeExecutorConfig: """Convert current instance to config object""" return JupyterCodeExecutorConfig( - kernel_name=self._kernel_name, timeout=self._timeout, output_dir=str(self._output_dir) + kernel_name=self._kernel_name, timeout=self._timeout, output_dir=str(self.output_dir) ) + @property + def output_dir(self) -> Path: + # If a user specifies the current directory, warn them that this is deprecated + if self._output_dir == Path("."): + warnings.warn( + "Using the current directory as output_dir is deprecated", + DeprecationWarning, + stacklevel=2, + ) + return self._output_dir + @classmethod def _from_config(cls, config: JupyterCodeExecutorConfig) -> Self: """Create instance from config object""" - return cls(kernel_name=config.kernel_name, timeout=config.timeout, output_dir=Path(config.output_dir)) + return cls( + kernel_name=config.kernel_name, + timeout=config.timeout, + output_dir=Path(config.output_dir) if config.output_dir else None, + ) diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/local/__init__.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/local/__init__.py index 852c4aee28d4..5ab7ecc91742 100644 --- a/python/packages/autogen-ext/src/autogen_ext/code_executors/local/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/local/__init__.py @@ -5,6 +5,7 @@ import logging import os import sys +import tempfile import warnings from hashlib import sha256 from pathlib import Path @@ -36,8 +37,9 @@ class LocalCommandLineCodeExecutorConfig(BaseModel): """Configuration for LocalCommandLineCodeExecutor""" timeout: int = 60 - work_dir: str = "." # Stored as string, converted to Path in _from_config + work_dir: Optional[str] = None functions_module: str = "functions" + cleanup_temp_files: bool = True class LocalCommandLineCodeExecutor(CodeExecutor, Component[LocalCommandLineCodeExecutorConfig]): @@ -56,7 +58,7 @@ class LocalCommandLineCodeExecutor(CodeExecutor, Component[LocalCommandLineCodeE commands from being executed which may potentially affect the users environment. Currently the only supported languages is Python and shell scripts. For Python code, use the language "python" for the code block. - For shell scripts, use the language "bash", "shell", or "sh" for the code + For shell scripts, use the language "bash", "shell", "sh", "pwsh", "powershell", or "ps1" for the code block. .. note:: @@ -74,12 +76,16 @@ class LocalCommandLineCodeExecutor(CodeExecutor, Component[LocalCommandLineCodeE Args: timeout (int): The timeout for the execution of any single code block. Default is 60. work_dir (str): The working directory for the code execution. If None, - a default working directory will be used. The default working - directory is the current directory ".". + a default working directory will be used. The default working directory is a temporary directory. functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]): A list of functions that are available to the code executor. Default is an empty list. functions_module (str, optional): The name of the module that will be created to store the functions. Defaults to "functions". + cleanup_temp_files (bool, optional): Whether to automatically clean up temporary files after execution. Defaults to True. virtual_env_context (Optional[SimpleNamespace], optional): The virtual environment context. Defaults to None. + .. note:: + Using the current directory (".") as working directory is deprecated. Using it will raise a deprecation warning. + + Example: How to use `LocalCommandLineCodeExecutor` with a virtual environment different from the one used to run the autogen application: @@ -141,7 +147,7 @@ async def example(): def __init__( self, timeout: int = 60, - work_dir: Union[Path, str] = Path("."), + work_dir: Optional[Union[Path, str]] = None, functions: Sequence[ Union[ FunctionWithRequirements[Any, A], @@ -150,23 +156,27 @@ def __init__( ] ] = [], functions_module: str = "functions", + cleanup_temp_files: bool = True, virtual_env_context: Optional[SimpleNamespace] = None, ): if timeout < 1: raise ValueError("Timeout must be greater than or equal to 1.") - - if isinstance(work_dir, str): - work_dir = Path(work_dir) - - if not functions_module.isidentifier(): - raise ValueError("Module name must be a valid Python identifier") - - self._functions_module = functions_module - - work_dir.mkdir(exist_ok=True) - self._timeout = timeout - self._work_dir: Path = work_dir + + self._work_dir: Optional[Path] = None + if work_dir is not None: + # Check if user provided work_dir is the current directory and warn if so. + if Path(work_dir).resolve() == Path.cwd().resolve(): + warnings.warn( + "Using the current directory as work_dir is deprecated.", + DeprecationWarning, + stacklevel=2, + ) + if isinstance(work_dir, str): + self._work_dir = Path(work_dir) + else: + self._work_dir = work_dir + self._work_dir.mkdir(exist_ok=True) self._functions = functions # Setup could take some time so we intentionally wait for the first code block to do it. @@ -175,8 +185,16 @@ def __init__( else: self._setup_functions_complete = True + if not functions_module.isidentifier(): + raise ValueError("Module name must be a valid Python identifier") + self._functions_module = functions_module + + self._cleanup_temp_files = cleanup_temp_files self._virtual_env_context: Optional[SimpleNamespace] = virtual_env_context + self._temp_dir: Optional[tempfile.TemporaryDirectory[str]] = None + self._started = False + # Check the current event loop policy if on windows. if sys.platform == "win32": current_policy = asyncio.get_event_loop_policy() @@ -213,27 +231,39 @@ def format_functions_for_prompt(self, prompt_template: str = FUNCTION_PROMPT_TEM ) @property - def functions_module(self) -> str: - """(Experimental) The module name for the functions.""" - return self._functions_module + def timeout(self) -> int: + """(Experimental) The timeout for code execution.""" + return self._timeout + + @property + def work_dir(self) -> Path: + """(Experimental) The working directory for the code execution.""" + if self._work_dir is not None: + return self._work_dir + else: + # Automatically create temp directory if not exists + if self._temp_dir is None: + self._temp_dir = tempfile.TemporaryDirectory() + self._started = True + return Path(self._temp_dir.name) @property def functions(self) -> List[str]: raise NotImplementedError @property - def timeout(self) -> int: - """(Experimental) The timeout for code execution.""" - return self._timeout + def functions_module(self) -> str: + """(Experimental) The module name for the functions.""" + return self._functions_module @property - def work_dir(self) -> Path: - """(Experimental) The working directory for the code execution.""" - return self._work_dir + def cleanup_temp_files(self) -> bool: + """(Experimental) Whether to automatically clean up temporary files after execution.""" + return self._cleanup_temp_files async def _setup_functions(self, cancellation_token: CancellationToken) -> None: func_file_content = build_python_functions_file(self._functions) - func_file = self._work_dir / f"{self._functions_module}.py" + func_file = self.work_dir / f"{self._functions_module}.py" func_file.write_text(func_file_content) # Collect requirements @@ -255,7 +285,7 @@ async def _setup_functions(self, cancellation_token: CancellationToken) -> None: asyncio.create_subprocess_exec( py_executable, *cmd_args, - cwd=self._work_dir, + cwd=self.work_dir, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) @@ -329,7 +359,7 @@ async def _execute_code_dont_check_setup( # Try extracting a filename (if present) try: - filename = get_file_name_from_content(code, self._work_dir) + filename = get_file_name_from_content(code, self.work_dir) except ValueError: return CommandLineCodeResult( exit_code=1, @@ -349,7 +379,7 @@ async def _execute_code_dont_check_setup( filename = f"tmp_code_{code_hash}.{ext}" - written_file = (self._work_dir / filename).resolve() + written_file = (self.work_dir / filename).resolve() with written_file.open("w", encoding="utf-8") as f: f.write(code) file_names.append(written_file) @@ -388,7 +418,7 @@ async def _execute_code_dont_check_setup( asyncio.create_subprocess_exec( program, *extra_args, - cwd=self._work_dir, + cwd=self.work_dir, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env=env, @@ -423,7 +453,16 @@ async def _execute_code_dont_check_setup( break code_file = str(file_names[0]) if file_names else None - return CommandLineCodeResult(exit_code=exitcode, output=logs_all, code_file=code_file) + code_result = CommandLineCodeResult(exit_code=exitcode, output=logs_all, code_file=code_file) + + if self._cleanup_temp_files: + for file in file_names: + try: + file.unlink(missing_ok=True) + except OSError as error: + logging.error(f"Failed to delete temporary file {file}: {error}") + + return code_result async def restart(self) -> None: """(Experimental) Restart the code executor.""" @@ -433,13 +472,26 @@ async def restart(self) -> None: ) async def start(self) -> None: - """(Experimental) Start the code executor.""" - # No action needed for local command line executor - pass + """(Experimental) Start the code executor. + + Initializes the local code executor and should be called before executing any code blocks. + It marks the executor internal state as started. + If no working directory is provided, the method creates a temporary directory for the executor to use. + """ + if self._work_dir is None and self._temp_dir is None: + self._temp_dir = tempfile.TemporaryDirectory() + self._started = True async def stop(self) -> None: - """(Experimental) Stop the code executor.""" - # No action needed for local command line executor + """(Experimental) Stop the code executor. + + Stops the local code executor and performs the cleanup of the temporary working directory (if it was created). + The executor's internal state is markes as no longer started. + """ + if self._temp_dir is not None: + self._temp_dir.cleanup() + self._temp_dir = None + self._started = False pass def _to_config(self) -> LocalCommandLineCodeExecutorConfig: @@ -450,14 +502,16 @@ def _to_config(self) -> LocalCommandLineCodeExecutorConfig: return LocalCommandLineCodeExecutorConfig( timeout=self._timeout, - work_dir=str(self._work_dir), + work_dir=str(self.work_dir), functions_module=self._functions_module, + cleanup_temp_files=self._cleanup_temp_files, ) @classmethod def _from_config(cls, config: LocalCommandLineCodeExecutorConfig) -> Self: return cls( timeout=config.timeout, - work_dir=Path(config.work_dir), + work_dir=Path(config.work_dir) if config.work_dir is not None else None, functions_module=config.functions_module, + cleanup_temp_files=config.cleanup_temp_files, ) diff --git a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/__init__.py b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/__init__.py index f68a3197d9fc..97415af2beda 100644 --- a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/__init__.py @@ -1,3 +1,4 @@ -from .memory_controller import MemoryController +from ._memory_bank import MemoryBankConfig +from .memory_controller import MemoryController, MemoryControllerConfig -__all__ = ["MemoryController"] +__all__ = ["MemoryController", "MemoryControllerConfig", "MemoryBankConfig"] diff --git a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/_prompter.py b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/_prompter.py index 2602feeb2c61..71bb4e7a5d44 100644 --- a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/_prompter.py +++ b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/_prompter.py @@ -184,7 +184,7 @@ async def find_index_topics(self, input_string: str) -> List[str]: return topic_list - async def generalize_task(self, task_description: str) -> str: + async def generalize_task(self, task_description: str, revise: bool | None = True) -> str: """ Attempts to rewrite a task description in a more general form. """ @@ -198,29 +198,31 @@ async def generalize_task(self, task_description: str) -> str: user_message.append(task_description) self._clear_history() - await self.call_model( + generalized_task = await self.call_model( summary="Ask the model to rephrase the task in a list of important points", system_message_content=sys_message, user_content=user_message, ) - user_message = [ - "Do you see any parts of this list that are irrelevant to actually solving the task? If so, explain which items are irrelevant." - ] - await self.call_model( - summary="Ask the model to identify irrelevant points", - system_message_content=sys_message, - user_content=user_message, - ) + if revise: + user_message = [ + "Do you see any parts of this list that are irrelevant to actually solving the task? If so, explain which items are irrelevant." + ] + await self.call_model( + summary="Ask the model to identify irrelevant points", + system_message_content=sys_message, + user_content=user_message, + ) + + user_message = [ + "Revise your original list to include only the most general terms, those that are critical to solving the task, removing any themes or descriptions that are not essential to the solution. Your final list may be shorter, but do not leave out any part of the task that is needed for solving the task. Do not add any additional commentary either before or after the list." + ] + generalized_task = await self.call_model( + summary="Ask the model to make a final list of general terms", + system_message_content=sys_message, + user_content=user_message, + ) - user_message = [ - "Revise your original list to include only the most general terms, those that are critical to solving the task, removing any themes or descriptions that are not essential to the solution. Your final list may be shorter, but do not leave out any part of the task that is needed for solving the task. Do not add any additional commentary either before or after the list." - ] - generalized_task = await self.call_model( - summary="Ask the model to make a final list of general terms", - system_message_content=sys_message, - user_content=user_message, - ) return generalized_task async def validate_insight(self, insight: str, task_description: str) -> bool: diff --git a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/memory_controller.py b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/memory_controller.py index 3a25f6ea18f9..acf5a649d72f 100644 --- a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/memory_controller.py +++ b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/memory_controller.py @@ -16,6 +16,11 @@ # Following the nested-config pattern, this TypedDict minimizes code changes by encapsulating # the settings that change frequently, as when loading many settings from a single YAML file. class MemoryControllerConfig(TypedDict, total=False): + generalize_task: bool + revise_generalized_task: bool + generate_topics: bool + validate_memos: bool + max_memos_to_retrieve: int max_train_trials: int max_test_trials: int MemoryBank: "MemoryBankConfig" @@ -33,6 +38,11 @@ class MemoryController: task_assignment_callback: An optional callback used to assign a task to any agent managed by the caller. config: An optional dict that can be used to override the following values: + - generalize_task: Whether to rewrite tasks in more general terms. + - revise_generalized_task: Whether to critique then rewrite the generalized task. + - generate_topics: Whether to base retrieval directly on tasks, or on topics extracted from tasks. + - validate_memos: Whether to apply a final validation stage to retrieved memos. + - max_memos_to_retrieve: The maximum number of memos to return from retrieve_relevant_memos(). - max_train_trials: The maximum number of learning iterations to attempt when training on a task. - max_test_trials: The total number of attempts made when testing for failure on a task. - MemoryBank: A config dict passed to MemoryBank. @@ -91,10 +101,20 @@ def __init__( self.logger.enter_function() # Apply default settings and any config overrides. + self.generalize_task = True + self.revise_generalized_task = True + self.generate_topics = True + self.validate_memos = True + self.max_memos_to_retrieve = 10 self.max_train_trials = 10 self.max_test_trials = 3 memory_bank_config = None if config is not None: + self.generalize_task = config.get("generalize_task", self.generalize_task) + self.revise_generalized_task = config.get("revise_generalized_task", self.revise_generalized_task) + self.generate_topics = config.get("generate_topics", self.generate_topics) + self.validate_memos = config.get("validate_memos", self.validate_memos) + self.max_memos_to_retrieve = config.get("max_memos_to_retrieve", self.max_memos_to_retrieve) self.max_train_trials = config.get("max_train_trials", self.max_train_trials) self.max_test_trials = config.get("max_test_trials", self.max_test_trials) memory_bank_config = config.get("MemoryBank", memory_bank_config) @@ -178,8 +198,10 @@ async def add_memo(self, insight: str, task: None | str = None, index_on_both: b if task is not None: self.logger.info("\nGIVEN TASK:") self.logger.info(task) - # Generalize the task. - generalized_task = await self.prompter.generalize_task(task) + if self.generalize_task: + generalized_task = await self.prompter.generalize_task(task, revise=self.revise_generalized_task) + else: + generalized_task = task self.logger.info("\nGIVEN INSIGHT:") self.logger.info(insight) @@ -196,7 +218,10 @@ async def add_memo(self, insight: str, task: None | str = None, index_on_both: b text_to_index = task self.logger.info("\nTOPICS EXTRACTED FROM TASK:") - topics = await self.prompter.find_index_topics(text_to_index) + if self.generate_topics: + topics = await self.prompter.find_index_topics(text_to_index) + else: + topics = [text_to_index] self.logger.info("\n".join(topics)) self.logger.info("") @@ -218,7 +243,10 @@ async def add_task_solution_pair_to_memory(self, task: str, solution: str) -> No self.logger.info(solution) # Get a list of topics from the task. - topics = await self.prompter.find_index_topics(task.strip()) + if self.generate_topics: + topics = await self.prompter.find_index_topics(task.strip()) + else: + topics = [task.strip()] self.logger.info("\nTOPICS EXTRACTED FROM TASK:") self.logger.info("\n".join(topics)) self.logger.info("") @@ -238,8 +266,14 @@ async def retrieve_relevant_memos(self, task: str) -> List[Memo]: self.logger.info(task) # Get a list of topics from the generalized task. - generalized_task = await self.prompter.generalize_task(task) - task_topics = await self.prompter.find_index_topics(generalized_task) + if self.generalize_task: + generalized_task = await self.prompter.generalize_task(task, revise=self.revise_generalized_task) + else: + generalized_task = task + if self.generate_topics: + task_topics = await self.prompter.find_index_topics(generalized_task) + else: + task_topics = [generalized_task] self.logger.info("\nTOPICS EXTRACTED FROM TASK:") self.logger.info("\n".join(task_topics)) self.logger.info("") @@ -250,7 +284,9 @@ async def retrieve_relevant_memos(self, task: str) -> List[Memo]: # Apply a final validation stage to keep only the memos that the LLM concludes are sufficiently relevant. validated_memos: List[Memo] = [] for memo in memo_list: - if await self.prompter.validate_insight(memo.insight, task): + if len(validated_memos) >= self.max_memos_to_retrieve: + break + if (not self.validate_memos) or await self.prompter.validate_insight(memo.insight, task): validated_memos.append(memo) self.logger.info("\n{} VALIDATED MEMOS".format(len(validated_memos))) diff --git a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/__init__.py b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/__init__.py index 992e0e8ef1c0..82bf516d28d0 100644 --- a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/__init__.py @@ -1,7 +1,15 @@ -from .apprentice import Apprentice +from .apprentice import Apprentice, ApprenticeConfig from .chat_completion_client_recorder import ChatCompletionClientRecorder from .grader import Grader -from .page_logger import PageLogger +from .page_logger import PageLogger, PageLoggerConfig from .teachability import Teachability -__all__ = ["Apprentice", "ChatCompletionClientRecorder", "Grader", "PageLogger", "Teachability"] +__all__ = [ + "Apprentice", + "ChatCompletionClientRecorder", + "Grader", + "PageLogger", + "Teachability", + "ApprenticeConfig", + "PageLoggerConfig", +] diff --git a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/apprentice.py b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/apprentice.py index a8104c0ebc44..b212628ecdbb 100644 --- a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/apprentice.py +++ b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/apprentice.py @@ -4,7 +4,7 @@ from autogen_agentchat.agents import AssistantAgent from autogen_agentchat.base import TaskResult -from autogen_agentchat.messages import AgentEvent, ChatMessage, TextMessage +from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage, TextMessage from autogen_core.models import ( ChatCompletionClient, LLMMessage, @@ -190,9 +190,9 @@ async def _assign_task_to_assistant_agent(self, task: str) -> Tuple[Any, Any]: # Get the agent's response to the task. task_result: TaskResult = await assistant_agent.run(task=TextMessage(content=task, source="User")) - messages: Sequence[AgentEvent | ChatMessage] = task_result.messages - message: AgentEvent | ChatMessage = messages[-1] - response_str = message.content + messages: Sequence[BaseAgentEvent | BaseChatMessage] = task_result.messages + message: BaseAgentEvent | BaseChatMessage = messages[-1] + response_str = message.to_text() # Log the model call self.logger.log_model_task( @@ -245,12 +245,7 @@ async def _assign_task_to_magentic_one(self, task: str) -> Tuple[str, str]: response_str_list: List[str] = [] for message in messages: - content = message.content - if isinstance(content, str): - content_str = content - else: - content_str = "Not a string." - response_str_list.append(content_str) + response_str_list.append(message.to_text()) response_str = "\n".join(response_str_list) self.logger.info("\n----- RESPONSE -----\n\n{}\n".format(response_str)) diff --git a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/chat_completion_client_recorder.py b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/chat_completion_client_recorder.py index 16124db1f3c5..8b981312f427 100644 --- a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/chat_completion_client_recorder.py +++ b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/chat_completion_client_recorder.py @@ -41,10 +41,9 @@ class ChatCompletionClientRecorder(ChatCompletionClient): create calls) or a "stream" (a list of streamed outputs for create_stream calls). ReplayChatCompletionClient and ChatCompletionCache do similar things, but with significant differences: - - ReplayChatCompletionClient replays pre-defined responses in a specified order - without recording anything or checking the messages sent to the client. - - ChatCompletionCache caches responses and replays them for messages that have been seen before, - regardless of order, and calls the base client for any uncached messages. + + - ReplayChatCompletionClient replays pre-defined responses in a specified order without recording anything or checking the messages sent to the client. + - ChatCompletionCache caches responses and replays them for messages that have been seen before, regardless of order, and calls the base client for any uncached messages. """ def __init__( @@ -91,6 +90,7 @@ async def create( json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", ) -> CreateResult: current_messages: List[Mapping[str, Any]] = [msg.model_dump() for msg in messages] if self.mode == "record": @@ -98,6 +98,7 @@ async def create( messages, tools=tools, json_output=json_output, + tool_choice=tool_choice, extra_create_args=extra_create_args, cancellation_token=cancellation_token, ) @@ -158,10 +159,12 @@ def create_stream( json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", ) -> AsyncGenerator[Union[str, CreateResult], None]: return self.base_client.create_stream( messages, tools=tools, + tool_choice=tool_choice, json_output=json_output, extra_create_args=extra_create_args, cancellation_token=cancellation_token, diff --git a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/page_logger.py b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/page_logger.py index 92964dfbec12..fa7fe2f1d567 100644 --- a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/page_logger.py +++ b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/page_logger.py @@ -5,7 +5,7 @@ from typing import Any, Dict, List, Mapping, Optional, Sequence, TypedDict from autogen_agentchat.base import TaskResult -from autogen_agentchat.messages import AgentEvent, ChatMessage +from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage from autogen_core import Image from autogen_core.models import ( AssistantMessage, @@ -343,9 +343,9 @@ def log_model_task( if self.level > self.levels["INFO"]: return None - messages: Sequence[AgentEvent | ChatMessage] = task_result.messages + messages: Sequence[BaseAgentEvent | BaseChatMessage] = task_result.messages message = messages[-1] - response_str = message.content + response_str = message.to_text() if not isinstance(response_str, str): response_str = "??" diff --git a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/teachability.py b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/teachability.py index f8a09ee40e34..d9f511b93201 100644 --- a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/teachability.py +++ b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/teachability.py @@ -14,10 +14,11 @@ class Teachability(Memory): Gives an AssistantAgent the ability to learn quickly from user teachings, hints, and advice. Steps for usage: - 1. Instantiate MemoryController. - 2. Instantiate Teachability, passing the memory controller as a parameter. - 3. Instantiate an AssistantAgent, passing the teachability instance (wrapped in a list) as the memory parameter. - 4. Use the AssistantAgent as usual, such as for chatting with the user. + + 1. Instantiate MemoryController. + 2. Instantiate Teachability, passing the memory controller as a parameter. + 3. Instantiate an AssistantAgent, passing the teachability instance (wrapped in a list) as the memory parameter. + 4. Use the AssistantAgent as usual, such as for chatting with the user. """ def __init__(self, memory_controller: "MemoryController", name: str | None = None) -> None: diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/canvas/__init__.py b/python/packages/autogen-ext/src/autogen_ext/memory/canvas/__init__.py new file mode 100644 index 000000000000..ad10924579fb --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/memory/canvas/__init__.py @@ -0,0 +1,4 @@ +from ._text_canvas import TextCanvas +from ._text_canvas_memory import TextCanvasMemory + +__all__ = ["TextCanvas", "TextCanvasMemory"] diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/canvas/_canvas.py b/python/packages/autogen-ext/src/autogen_ext/memory/canvas/_canvas.py new file mode 100644 index 000000000000..de2eca26b9ca --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/memory/canvas/_canvas.py @@ -0,0 +1,50 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict, Union + + +class BaseCanvas(ABC): + """ + An abstract protocol for "canvas" objects that maintain + revision history for file-like data. Concrete subclasses + can handle text, images, structured data, etc. + + .. warning:: + + This is an experimental API and may change in the future. + + """ + + @abstractmethod + def list_files(self) -> Dict[str, int]: + """ + Returns a dict of filename -> latest revision number. + """ + raise NotImplementedError + + @abstractmethod + def get_latest_content(self, filename: str) -> Union[str, bytes, Any]: + """ + Returns the latest version of a file's content. + """ + raise NotImplementedError + + @abstractmethod + def add_or_update_file(self, filename: str, new_content: Union[str, bytes, Any]) -> None: + """ + Creates or updates the file content with a new revision. + """ + raise NotImplementedError + + @abstractmethod + def get_diff(self, filename: str, from_revision: int, to_revision: int) -> str: + """ + Returns a diff (in some format) between two revisions. + """ + raise NotImplementedError + + @abstractmethod + def apply_patch(self, filename: str, patch_data: Union[str, bytes, Any]) -> None: + """ + Applies a patch/diff to the latest revision and increments the revision. + """ + raise NotImplementedError diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/canvas/_canvas_writer.py b/python/packages/autogen-ext/src/autogen_ext/memory/canvas/_canvas_writer.py new file mode 100644 index 000000000000..b125eb9ce075 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/memory/canvas/_canvas_writer.py @@ -0,0 +1,64 @@ +from autogen_core import CancellationToken +from autogen_core.tools import BaseTool +from pydantic import BaseModel + +from ._text_canvas import TextCanvas + + +class UpdateFileArgs(BaseModel): + filename: str + new_content: str + + +class UpdateFileResult(BaseModel): + status: str + + +class UpdateFileTool(BaseTool[UpdateFileArgs, UpdateFileResult]): + """ + Overwrites or creates a file in the canvas. + """ + + def __init__(self, canvas: TextCanvas): + super().__init__( + args_type=UpdateFileArgs, + return_type=UpdateFileResult, + name="update_file", + description="Create/update a file on the canvas with the provided content.", + ) + self._canvas = canvas + + async def run(self, args: UpdateFileArgs, cancellation_token: CancellationToken) -> UpdateFileResult: + self._canvas.add_or_update_file(args.filename, args.new_content) + return UpdateFileResult(status="OK") + + +class ApplyPatchArgs(BaseModel): + filename: str + patch_text: str + + +class ApplyPatchResult(BaseModel): + status: str + + +class ApplyPatchTool(BaseTool[ApplyPatchArgs, ApplyPatchResult]): + """ + Applies a unified diff patch to the given file on the canvas. + """ + + def __init__(self, canvas: TextCanvas): + super().__init__( + args_type=ApplyPatchArgs, + return_type=ApplyPatchResult, + name="apply_patch", + description=( + "Apply a unified diff patch to an existing file on the canvas. " + "The patch must be in diff/patch format. The file must exist or be created first." + ), + ) + self._canvas = canvas + + async def run(self, args: ApplyPatchArgs, cancellation_token: CancellationToken) -> ApplyPatchResult: + self._canvas.apply_patch(args.filename, args.patch_text) + return ApplyPatchResult(status="PATCH APPLIED") diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/canvas/_text_canvas.py b/python/packages/autogen-ext/src/autogen_ext/memory/canvas/_text_canvas.py new file mode 100644 index 000000000000..306a070a4299 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/memory/canvas/_text_canvas.py @@ -0,0 +1,192 @@ +import difflib +from typing import Any, Dict, List, Union + +try: # pragma: no cover + from unidiff import PatchSet +except ModuleNotFoundError: # pragma: no cover + PatchSet = None # type: ignore + +from ._canvas import BaseCanvas + + +class FileRevision: + """Tracks the history of one file's content.""" + + __slots__ = ("content", "revision") + + def __init__(self, content: str, revision: int) -> None: + self.content: str = content + self.revision: int = revision # e.g. an integer, a timestamp, or git hash + + +class TextCanvas(BaseCanvas): + """An in‑memory canvas that stores *text* files with full revision history. + + .. warning:: + + This is an experimental API and may change in the future. + + Besides the original CRUD‑like operations, this enhanced implementation adds: + + * **apply_patch** – applies patches using the ``unidiff`` library for accurate + hunk application and context line validation. + * **get_revision_content** – random access to any historical revision. + * **get_revision_diffs** – obtain the list of diffs applied between every + consecutive pair of revisions so that a caller can replay or audit the + full change history. + """ + + # ---------------------------------------------------------------------------------- + # Construction helpers + # ---------------------------------------------------------------------------------- + + def __init__(self) -> None: + # For each file we keep an *ordered* list of FileRevision where the last + # element is the most recent. Using a list keeps the memory footprint + # small and preserves order without any extra bookkeeping. + self._files: Dict[str, List[FileRevision]] = {} + + # ---------------------------------------------------------------------------------- + # Internal utilities + # ---------------------------------------------------------------------------------- + + def _latest_idx(self, filename: str) -> int: + """Return the index (not revision number) of the newest revision.""" + return len(self._files.get(filename, [])) - 1 + + def _ensure_file(self, filename: str) -> None: + if filename not in self._files: + raise ValueError(f"File '{filename}' does not exist on the canvas; create it first.") + + # ---------------------------------------------------------------------------------- + # Revision inspection helpers + # ---------------------------------------------------------------------------------- + + def get_revision_content(self, filename: str, revision: int) -> str: # NEW 🚀 + """Return the exact content stored in *revision*. + + If the revision does not exist an empty string is returned so that + downstream code can handle the "not found" case without exceptions. + """ + for rev in self._files.get(filename, []): + if rev.revision == revision: + return rev.content + return "" + + def get_revision_diffs(self, filename: str) -> List[str]: # NEW 🚀 + """Return a *chronological* list of unified‑diffs for *filename*. + + Each element in the returned list represents the diff that transformed + revision *n* into revision *n+1* (starting at revision 1 → 2). + """ + revisions = self._files.get(filename, []) + diffs: List[str] = [] + for i in range(1, len(revisions)): + older, newer = revisions[i - 1], revisions[i] + diff = difflib.unified_diff( + older.content.splitlines(keepends=True), + newer.content.splitlines(keepends=True), + fromfile=f"{filename}@r{older.revision}", + tofile=f"{filename}@r{newer.revision}", + ) + diffs.append("".join(diff)) + return diffs + + # ---------------------------------------------------------------------------------- + # BaseCanvas interface implementation + # ---------------------------------------------------------------------------------- + + def list_files(self) -> Dict[str, int]: + """Return a mapping of *filename → latest revision number*.""" + return {fname: revs[-1].revision for fname, revs in self._files.items() if revs} + + def get_latest_content(self, filename: str) -> str: # noqa: D401 – keep API identical + """Return the most recent content or an empty string if the file is new.""" + revs = self._files.get(filename, []) + return revs[-1].content if revs else "" + + def add_or_update_file(self, filename: str, new_content: Union[str, bytes, Any]) -> None: + """Create *filename* or append a new revision containing *new_content*.""" + if isinstance(new_content, bytes): + new_content = new_content.decode("utf-8") + if not isinstance(new_content, str): + raise ValueError(f"Expected str or bytes, got {type(new_content)}") + if filename not in self._files: + self._files[filename] = [FileRevision(new_content, 1)] + else: + last_rev_num = self._files[filename][-1].revision + self._files[filename].append(FileRevision(new_content, last_rev_num + 1)) + + def get_diff(self, filename: str, from_revision: int, to_revision: int) -> str: + """Return a unified diff between *from_revision* and *to_revision*.""" + revisions = self._files.get(filename, []) + if not revisions: + return "" + # Fetch the contents for the requested revisions. + from_content = self.get_revision_content(filename, from_revision) + to_content = self.get_revision_content(filename, to_revision) + if from_content == "" and to_content == "": # one (or both) revision ids not found + return "" + diff = difflib.unified_diff( + from_content.splitlines(keepends=True), + to_content.splitlines(keepends=True), + fromfile=f"{filename}@r{from_revision}", + tofile=f"{filename}@r{to_revision}", + ) + return "".join(diff) + + def apply_patch(self, filename: str, patch_data: Union[str, bytes, Any]) -> None: + """Apply *patch_text* (unified diff) to the latest revision and save a new revision. + + Uses the *unidiff* library to accurately apply hunks and validate context lines. + """ + if isinstance(patch_data, bytes): + patch_data = patch_data.decode("utf-8") + if not isinstance(patch_data, str): + raise ValueError(f"Expected str or bytes, got {type(patch_data)}") + self._ensure_file(filename) + original_content = self.get_latest_content(filename) + + if PatchSet is None: + raise ImportError( + "The 'unidiff' package is required for patch application. Install with 'pip install unidiff'." + ) + + patch = PatchSet(patch_data) + # Our canvas stores exactly one file per patch operation so we + # use the first (and only) patched_file object. + if not patch: + raise ValueError("Empty patch text provided.") + patched_file = patch[0] + working_lines = original_content.splitlines(keepends=True) + line_offset = 0 + for hunk in patched_file: + # Calculate the slice boundaries in the *current* working copy. + start = hunk.source_start - 1 + line_offset + end = start + hunk.source_length + # Build the replacement block for this hunk. + replacement: List[str] = [] + for line in hunk: + if line.is_added or line.is_context: + replacement.append(line.value) + # removed lines (line.is_removed) are *not* added. + # Replace the slice with the hunk‑result. + working_lines[start:end] = replacement + line_offset += len(replacement) - (end - start) + new_content = "".join(working_lines) + + # Finally commit the new revision. + self.add_or_update_file(filename, new_content) + + # ---------------------------------------------------------------------------------- + # Convenience helpers + # ---------------------------------------------------------------------------------- + + def get_all_contents_for_context(self) -> str: # noqa: D401 – keep public API stable + """Return a summarised view of every file and its *latest* revision.""" + out: List[str] = ["=== CANVAS FILES ==="] + for fname, revs in self._files.items(): + latest = revs[-1] + out.append(f"File: {fname} (rev {latest.revision}):\n{latest.content}\n") + out.append("=== END OF CANVAS ===") + return "\n".join(out) diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/canvas/_text_canvas_memory.py b/python/packages/autogen-ext/src/autogen_ext/memory/canvas/_text_canvas_memory.py new file mode 100644 index 000000000000..ecf41a3c9e95 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/memory/canvas/_text_canvas_memory.py @@ -0,0 +1,229 @@ +from typing import Any, Optional + +from autogen_core import CancellationToken +from autogen_core.memory import ( + Memory, + MemoryContent, + MemoryMimeType, + MemoryQueryResult, + UpdateContextResult, +) +from autogen_core.model_context import ChatCompletionContext +from autogen_core.models import SystemMessage + +from ._canvas_writer import ApplyPatchTool, UpdateFileTool +from ._text_canvas import TextCanvas + + +class TextCanvasMemory(Memory): + """ + A memory implementation that uses a Canvas for storing file-like content. + Inserts the current state of the canvas into the ChatCompletionContext on each turn. + + .. warning:: + + This is an experimental API and may change in the future. + + The TextCanvasMemory provides a persistent, file-like storage mechanism that can be used + by agents to read and write content. It automatically injects the current state of all files + in the canvas into the model context before each inference. + + This is particularly useful for: + - Allowing agents to create and modify documents over multiple turns + - Enabling collaborative document editing between multiple agents + - Maintaining persistent state across conversation turns + - Working with content too large to fit in a single message + + The canvas provides tools for: + - Creating or updating files with new content + - Applying patches (unified diff format) to existing files + + Examples: + + **Example: Using TextCanvasMemory with an AssistantAgent** + + The following example demonstrates how to create a TextCanvasMemory and use it with + an AssistantAgent to write and update a story file. + + .. code-block:: python + + import asyncio + from autogen_core import CancellationToken + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.messages import TextMessage + from autogen_ext.memory.canvas import TextCanvasMemory + + + async def main(): + # Create a model client + model_client = OpenAIChatCompletionClient( + model="gpt-4o", + # api_key = "your_openai_api_key" + ) + + # Create the canvas memory + text_canvas_memory = TextCanvasMemory() + + # Get tools for working with the canvas + update_file_tool = text_canvas_memory.get_update_file_tool() + apply_patch_tool = text_canvas_memory.get_apply_patch_tool() + + # Create an agent with the canvas memory and tools + writer_agent = AssistantAgent( + name="Writer", + model_client=model_client, + description="A writer agent that creates and updates stories.", + system_message=''' + You are a Writer Agent. Your focus is to generate a story based on the user's request. + + Instructions for using the canvas: + + - The story should be stored on the canvas in a file named "story.md". + - If "story.md" does not exist, create it by calling the 'update_file' tool. + - If "story.md" already exists, generate a unified diff (patch) from the current + content to the new version, and call the 'apply_patch' tool to apply the changes. + + IMPORTANT: Do not include the full story text in your chat messages. + Only write the story content to the canvas using the tools. + ''', + tools=[update_file_tool, apply_patch_tool], + memory=[text_canvas_memory], + ) + + # Send a message to the agent + await writer_agent.on_messages( + [TextMessage(content="Write a short story about a bunny and a sunflower.", source="user")], + CancellationToken(), + ) + + # Retrieve the content from the canvas + story_content = text_canvas_memory.canvas.get_latest_content("story.md") + print("Story content from canvas:") + print(story_content) + + + if __name__ == "__main__": + asyncio.run(main()) + + **Example: Using TextCanvasMemory with multiple agents** + + The following example shows how to use TextCanvasMemory with multiple agents + collaborating on the same document. + + .. code-block:: python + + import asyncio + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.teams import RoundRobinGroupChat + from autogen_agentchat.conditions import TextMentionTermination + from autogen_ext.memory.canvas import TextCanvasMemory + + + async def main(): + # Create a model client + model_client = OpenAIChatCompletionClient( + model="gpt-4o", + # api_key = "your_openai_api_key" + ) + + # Create the shared canvas memory + text_canvas_memory = TextCanvasMemory() + update_file_tool = text_canvas_memory.get_update_file_tool() + apply_patch_tool = text_canvas_memory.get_apply_patch_tool() + + # Create a writer agent + writer_agent = AssistantAgent( + name="Writer", + model_client=model_client, + description="A writer agent that creates stories.", + system_message="You write children's stories on the canvas in story.md.", + tools=[update_file_tool, apply_patch_tool], + memory=[text_canvas_memory], + ) + + # Create a critique agent + critique_agent = AssistantAgent( + name="Critique", + model_client=model_client, + description="A critique agent that provides feedback on stories.", + system_message="You review the story.md file and provide constructive feedback.", + memory=[text_canvas_memory], + ) + + # Create a team with both agents + team = RoundRobinGroupChat( + participants=[writer_agent, critique_agent], + termination_condition=TextMentionTermination("TERMINATE"), + max_turns=10, + ) + + # Run the team on a task + await team.run(task="Create a children's book about a bunny and a sunflower") + + # Get the final story + story = text_canvas_memory.canvas.get_latest_content("story.md") + print(story) + + + if __name__ == "__main__": + asyncio.run(main()) + """ + + def __init__(self, canvas: Optional[TextCanvas] = None): + super().__init__() + self.canvas = canvas if canvas is not None else TextCanvas() + + async def update_context(self, model_context: ChatCompletionContext) -> UpdateContextResult: + """ + Inject the entire canvas summary (or a selected subset) as reference data. + Here, we just put it into a system message, but you could customize. + """ + snapshot = self.canvas.get_all_contents_for_context() + if snapshot.strip(): + msg = SystemMessage(content=snapshot) + await model_context.add_message(msg) + + # Return it for debugging/logging + memory_content = MemoryContent(content=snapshot, mime_type=MemoryMimeType.TEXT) + return UpdateContextResult(memories=MemoryQueryResult(results=[memory_content])) + + return UpdateContextResult(memories=MemoryQueryResult(results=[])) + + async def query( + self, query: str | MemoryContent, cancellation_token: Optional[CancellationToken] = None, **kwargs: Any + ) -> MemoryQueryResult: + """ + Potentially search for matching filenames or file content. + This example returns empty. + """ + return MemoryQueryResult(results=[]) + + async def add(self, content: MemoryContent, cancellation_token: Optional[CancellationToken] = None) -> None: + """ + Example usage: Possibly interpret content as a patch or direct file update. + Could also be done by a specialized "CanvasTool" instead. + """ + # NO-OP here, leaving actual changes to the CanvasTool + pass + + async def clear(self) -> None: + """Clear the entire canvas by replacing it with a new empty instance.""" + # Create a new TextCanvas instance instead of calling __init__ directly + self.canvas = TextCanvas() + + async def close(self) -> None: + pass + + def get_update_file_tool(self) -> UpdateFileTool: + """ + Returns an UpdateFileTool instance that works with this memory's canvas. + """ + return UpdateFileTool(self.canvas) + + def get_apply_patch_tool(self) -> ApplyPatchTool: + """ + Returns an ApplyPatchTool instance that works with this memory's canvas. + """ + return ApplyPatchTool(self.canvas) diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/__init__.py b/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/__init__.py new file mode 100644 index 000000000000..1d6ad04a0285 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/__init__.py @@ -0,0 +1,21 @@ +from ._chroma_configs import ( + ChromaDBVectorMemoryConfig, + CustomEmbeddingFunctionConfig, + DefaultEmbeddingFunctionConfig, + HttpChromaDBVectorMemoryConfig, + OpenAIEmbeddingFunctionConfig, + PersistentChromaDBVectorMemoryConfig, + SentenceTransformerEmbeddingFunctionConfig, +) +from ._chromadb import ChromaDBVectorMemory + +__all__ = [ + "ChromaDBVectorMemory", + "ChromaDBVectorMemoryConfig", + "PersistentChromaDBVectorMemoryConfig", + "HttpChromaDBVectorMemoryConfig", + "DefaultEmbeddingFunctionConfig", + "SentenceTransformerEmbeddingFunctionConfig", + "OpenAIEmbeddingFunctionConfig", + "CustomEmbeddingFunctionConfig", +] diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chroma_configs.py b/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chroma_configs.py new file mode 100644 index 000000000000..3f30caacdbde --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chroma_configs.py @@ -0,0 +1,137 @@ +"""Configuration classes for ChromaDB vector memory.""" + +from typing import Any, Callable, Dict, Literal, Union + +from pydantic import BaseModel, Field +from typing_extensions import Annotated + + +class DefaultEmbeddingFunctionConfig(BaseModel): + """Configuration for the default ChromaDB embedding function. + + Uses ChromaDB's default embedding function (Sentence Transformers all-MiniLM-L6-v2). + + .. versionadded:: v0.4.1 + Support for custom embedding functions in ChromaDB memory. + """ + + function_type: Literal["default"] = "default" + + +class SentenceTransformerEmbeddingFunctionConfig(BaseModel): + """Configuration for SentenceTransformer embedding functions. + + Allows specifying a custom SentenceTransformer model for embeddings. + + .. versionadded:: v0.4.1 + Support for custom embedding functions in ChromaDB memory. + + Args: + model_name (str): Name of the SentenceTransformer model to use. + Defaults to "all-MiniLM-L6-v2". + + Example: + .. code-block:: python + + from autogen_ext.memory.chromadb import SentenceTransformerEmbeddingFunctionConfig + + _ = SentenceTransformerEmbeddingFunctionConfig(model_name="paraphrase-multilingual-mpnet-base-v2") + """ + + function_type: Literal["sentence_transformer"] = "sentence_transformer" + model_name: str = Field(default="all-MiniLM-L6-v2", description="SentenceTransformer model name to use") + + +class OpenAIEmbeddingFunctionConfig(BaseModel): + """Configuration for OpenAI embedding functions. + + Uses OpenAI's embedding API for generating embeddings. + + .. versionadded:: v0.4.1 + Support for custom embedding functions in ChromaDB memory. + + Args: + api_key (str): OpenAI API key. If empty, will attempt to use environment variable. + model_name (str): OpenAI embedding model name. Defaults to "text-embedding-ada-002". + + Example: + .. code-block:: python + + from autogen_ext.memory.chromadb import OpenAIEmbeddingFunctionConfig + + _ = OpenAIEmbeddingFunctionConfig(api_key="sk-...", model_name="text-embedding-3-small") + """ + + function_type: Literal["openai"] = "openai" + api_key: str = Field(default="", description="OpenAI API key") + model_name: str = Field(default="text-embedding-ada-002", description="OpenAI embedding model name") + + +class CustomEmbeddingFunctionConfig(BaseModel): + """Configuration for custom embedding functions. + + Allows using a custom function that returns a ChromaDB-compatible embedding function. + + .. versionadded:: v0.4.1 + Support for custom embedding functions in ChromaDB memory. + + .. warning:: + Configurations containing custom functions are not serializable. + + Args: + function (Callable): Function that returns a ChromaDB-compatible embedding function. + params (Dict[str, Any]): Parameters to pass to the function. + """ + + function_type: Literal["custom"] = "custom" + function: Callable[..., Any] = Field(description="Function that returns an embedding function") + params: Dict[str, Any] = Field(default_factory=dict, description="Parameters to pass to the function") + + +# Tagged union type for embedding function configurations +EmbeddingFunctionConfig = Annotated[ + Union[ + DefaultEmbeddingFunctionConfig, + SentenceTransformerEmbeddingFunctionConfig, + OpenAIEmbeddingFunctionConfig, + CustomEmbeddingFunctionConfig, + ], + Field(discriminator="function_type"), +] + + +class ChromaDBVectorMemoryConfig(BaseModel): + """Base configuration for ChromaDB-based memory implementation. + + .. versionchanged:: v0.4.1 + Added support for custom embedding functions via embedding_function_config. + """ + + client_type: Literal["persistent", "http"] + collection_name: str = Field(default="memory_store", description="Name of the ChromaDB collection") + distance_metric: str = Field(default="cosine", description="Distance metric for similarity search") + k: int = Field(default=3, description="Number of results to return in queries") + score_threshold: float | None = Field(default=None, description="Minimum similarity score threshold") + allow_reset: bool = Field(default=False, description="Whether to allow resetting the ChromaDB client") + tenant: str = Field(default="default_tenant", description="Tenant to use") + database: str = Field(default="default_database", description="Database to use") + embedding_function_config: EmbeddingFunctionConfig = Field( + default_factory=DefaultEmbeddingFunctionConfig, description="Configuration for the embedding function" + ) + + +class PersistentChromaDBVectorMemoryConfig(ChromaDBVectorMemoryConfig): + """Configuration for persistent ChromaDB memory.""" + + client_type: Literal["persistent", "http"] = "persistent" + persistence_path: str = Field(default="./chroma_db", description="Path for persistent storage") + + +class HttpChromaDBVectorMemoryConfig(ChromaDBVectorMemoryConfig): + """Configuration for HTTP ChromaDB memory.""" + + client_type: Literal["persistent", "http"] = "http" + host: str = Field(default="localhost", description="Host of the remote server") + port: int = Field(default=8000, description="Port of the remote server") + ssl: bool = Field(default=False, description="Whether to use HTTPS") + headers: Dict[str, str] | None = Field(default=None, description="Headers to send to the server") diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/chromadb.py b/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chromadb.py similarity index 55% rename from python/packages/autogen-ext/src/autogen_ext/memory/chromadb.py rename to python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chromadb.py index 35d8564098e1..6664b404407c 100644 --- a/python/packages/autogen-ext/src/autogen_ext/memory/chromadb.py +++ b/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chromadb.py @@ -1,18 +1,26 @@ import logging import uuid -from typing import Any, Dict, List, Literal +from typing import Any, List from autogen_core import CancellationToken, Component, Image from autogen_core.memory import Memory, MemoryContent, MemoryMimeType, MemoryQueryResult, UpdateContextResult from autogen_core.model_context import ChatCompletionContext from autogen_core.models import SystemMessage from chromadb import HttpClient, PersistentClient -from chromadb.api import ClientAPI from chromadb.api.models.Collection import Collection -from chromadb.api.types import Document, IncludeEnum, Metadata -from pydantic import BaseModel, Field +from chromadb.api.types import Document, Metadata from typing_extensions import Self +from ._chroma_configs import ( + ChromaDBVectorMemoryConfig, + CustomEmbeddingFunctionConfig, + DefaultEmbeddingFunctionConfig, + HttpChromaDBVectorMemoryConfig, + OpenAIEmbeddingFunctionConfig, + PersistentChromaDBVectorMemoryConfig, + SentenceTransformerEmbeddingFunctionConfig, +) + logger = logging.getLogger(__name__) @@ -24,36 +32,6 @@ ) from e -class ChromaDBVectorMemoryConfig(BaseModel): - """Base configuration for ChromaDB-based memory implementation.""" - - client_type: Literal["persistent", "http"] - collection_name: str = Field(default="memory_store", description="Name of the ChromaDB collection") - distance_metric: str = Field(default="cosine", description="Distance metric for similarity search") - k: int = Field(default=3, description="Number of results to return in queries") - score_threshold: float | None = Field(default=None, description="Minimum similarity score threshold") - allow_reset: bool = Field(default=False, description="Whether to allow resetting the ChromaDB client") - tenant: str = Field(default="default_tenant", description="Tenant to use") - database: str = Field(default="default_database", description="Database to use") - - -class PersistentChromaDBVectorMemoryConfig(ChromaDBVectorMemoryConfig): - """Configuration for persistent ChromaDB memory.""" - - client_type: Literal["persistent", "http"] = "persistent" - persistence_path: str = Field(default="./chroma_db", description="Path for persistent storage") - - -class HttpChromaDBVectorMemoryConfig(ChromaDBVectorMemoryConfig): - """Configuration for HTTP ChromaDB memory.""" - - client_type: Literal["persistent", "http"] = "http" - host: str = Field(default="localhost", description="Host of the remote server") - port: int = Field(default=8000, description="Port of the remote server") - ssl: bool = Field(default=False, description="Whether to use HTTPS") - headers: Dict[str, str] | None = Field(default=None, description="Headers to send to the server") - - class ChromaDBVectorMemory(Memory, Component[ChromaDBVectorMemoryConfig]): """ Store and retrieve memory using vector similarity search powered by ChromaDB. @@ -67,69 +45,138 @@ class ChromaDBVectorMemory(Memory, Component[ChromaDBVectorMemoryConfig]): For advanced use cases requiring specialized formatting of retrieved content, users should extend this class and override the `update_context()` method. - .. note:: + This implementation requires the ChromaDB extra to be installed. Install with: - This implementation requires the ChromaDB extra to be installed. Install with: - `pip install autogen-ext[chromadb]` + .. code-block:: bash + + pip install "autogen-ext[chromadb]" Args: config (ChromaDBVectorMemoryConfig | None): Configuration for the ChromaDB memory. If None, defaults to a PersistentChromaDBVectorMemoryConfig with default values. Two config types are supported: - - PersistentChromaDBVectorMemoryConfig: For local storage - - HttpChromaDBVectorMemoryConfig: For connecting to a remote ChromaDB server + * PersistentChromaDBVectorMemoryConfig: For local storage + * HttpChromaDBVectorMemoryConfig: For connecting to a remote ChromaDB server Example: .. code-block:: python import os + import asyncio from pathlib import Path from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.ui import Console from autogen_core.memory import MemoryContent, MemoryMimeType - from autogen_ext.memory.chromadb import ChromaDBVectorMemory, PersistentChromaDBVectorMemoryConfig + from autogen_ext.memory.chromadb import ( + ChromaDBVectorMemory, + PersistentChromaDBVectorMemoryConfig, + SentenceTransformerEmbeddingFunctionConfig, + OpenAIEmbeddingFunctionConfig, + ) from autogen_ext.models.openai import OpenAIChatCompletionClient - # Initialize ChromaDB memory with custom config - memory = ChromaDBVectorMemory( - config=PersistentChromaDBVectorMemoryConfig( - collection_name="user_preferences", - persistence_path=os.path.join(str(Path.home()), ".chromadb_autogen"), - k=3, # Return top 3 results - score_threshold=0.5, # Minimum similarity score + + def get_weather(city: str) -> str: + return f"The weather in {city} is sunny with a high of 90°F and a low of 70°F." + + + def fahrenheit_to_celsius(fahrenheit: float) -> float: + return (fahrenheit - 32) * 5.0 / 9.0 + + + async def main() -> None: + # Use default embedding function + default_memory = ChromaDBVectorMemory( + config=PersistentChromaDBVectorMemoryConfig( + collection_name="user_preferences", + persistence_path=os.path.join(str(Path.home()), ".chromadb_autogen"), + k=3, # Return top 3 results + score_threshold=0.5, # Minimum similarity score + ) ) - ) - # Add user preferences to memory - await memory.add( - MemoryContent( - content="The user prefers temperatures in Celsius", - mime_type=MemoryMimeType.TEXT, - metadata={"category": "preferences", "type": "units"}, + # Using a custom SentenceTransformer model + custom_memory = ChromaDBVectorMemory( + config=PersistentChromaDBVectorMemoryConfig( + collection_name="multilingual_memory", + persistence_path=os.path.join(str(Path.home()), ".chromadb_autogen"), + embedding_function_config=SentenceTransformerEmbeddingFunctionConfig( + model_name="paraphrase-multilingual-mpnet-base-v2" + ), + ) ) - ) - # Create assistant agent with ChromaDB memory - assistant = AssistantAgent( - name="assistant", - model_client=OpenAIChatCompletionClient( - model="gpt-4o", - ), - memory=[memory], - ) + # Using OpenAI embeddings + openai_memory = ChromaDBVectorMemory( + config=PersistentChromaDBVectorMemoryConfig( + collection_name="openai_memory", + persistence_path=os.path.join(str(Path.home()), ".chromadb_autogen"), + embedding_function_config=OpenAIEmbeddingFunctionConfig( + api_key=os.environ["OPENAI_API_KEY"], model_name="text-embedding-3-small" + ), + ) + ) + + # Add user preferences to memory + await openai_memory.add( + MemoryContent( + content="The user prefers weather temperatures in Celsius", + mime_type=MemoryMimeType.TEXT, + metadata={"category": "preferences", "type": "units"}, + ) + ) + + # Create assistant agent with ChromaDB memory + assistant = AssistantAgent( + name="assistant", + model_client=OpenAIChatCompletionClient( + model="gpt-4.1", + ), + tools=[ + get_weather, + fahrenheit_to_celsius, + ], + max_tool_iterations=10, + memory=[openai_memory], + ) + + # The memory will automatically retrieve relevant content during conversations + await Console(assistant.run_stream(task="What's the temperature in New York?")) + + # Remember to close the memory when finished + await default_memory.close() + await custom_memory.close() + await openai_memory.close() + - # The memory will automatically retrieve relevant content during conversations - stream = assistant.run_stream(task="What's the weather in New York?") + asyncio.run(main()) + + Output: + + .. code-block:: text + + ---------- TextMessage (user) ---------- + What's the temperature in New York? + ---------- MemoryQueryEvent (assistant) ---------- + [MemoryContent(content='The user prefers weather temperatures in Celsius', mime_type='MemoryMimeType.TEXT', metadata={'type': 'units', 'category': 'preferences', 'mime_type': 'MemoryMimeType.TEXT', 'score': 0.3133561611175537, 'id': 'fb00506c-acf4-4174-93d7-2a942593f3f7'}), MemoryContent(content='The user prefers weather temperatures in Celsius', mime_type='MemoryMimeType.TEXT', metadata={'mime_type': 'MemoryMimeType.TEXT', 'category': 'preferences', 'type': 'units', 'score': 0.3133561611175537, 'id': '34311689-b419-4e1a-8bc4-09143f356c66'})] + ---------- ToolCallRequestEvent (assistant) ---------- + [FunctionCall(id='call_7TjsFd430J1aKwU5T2w8bvdh', arguments='{"city":"New York"}', name='get_weather')] + ---------- ToolCallExecutionEvent (assistant) ---------- + [FunctionExecutionResult(content='The weather in New York is sunny with a high of 90°F and a low of 70°F.', name='get_weather', call_id='call_7TjsFd430J1aKwU5T2w8bvdh', is_error=False)] + ---------- ToolCallRequestEvent (assistant) ---------- + [FunctionCall(id='call_RTjMHEZwDXtjurEYTjDlvq9c', arguments='{"fahrenheit": 90}', name='fahrenheit_to_celsius'), FunctionCall(id='call_3mMuCK1aqtzZPTqIHPoHKxtP', arguments='{"fahrenheit": 70}', name='fahrenheit_to_celsius')] + ---------- ToolCallExecutionEvent (assistant) ---------- + [FunctionExecutionResult(content='32.22222222222222', name='fahrenheit_to_celsius', call_id='call_RTjMHEZwDXtjurEYTjDlvq9c', is_error=False), FunctionExecutionResult(content='21.11111111111111', name='fahrenheit_to_celsius', call_id='call_3mMuCK1aqtzZPTqIHPoHKxtP', is_error=False)] + ---------- TextMessage (assistant) ---------- + The temperature in New York today is sunny with a high of about 32°C and a low of about 21°C. - # Remember to close the memory when finished - await memory.close() """ component_config_schema = ChromaDBVectorMemoryConfig component_provider_override = "autogen_ext.memory.chromadb.ChromaDBVectorMemory" def __init__(self, config: ChromaDBVectorMemoryConfig | None = None) -> None: - """Initialize ChromaDBVectorMemory.""" self._config = config or PersistentChromaDBVectorMemoryConfig() self._client: ClientAPI | None = None self._collection: Collection | None = None @@ -139,6 +186,55 @@ def collection_name(self) -> str: """Get the name of the ChromaDB collection.""" return self._config.collection_name + def _create_embedding_function(self) -> Any: + """Create an embedding function based on the configuration. + + Returns: + A ChromaDB-compatible embedding function. + + Raises: + ValueError: If the embedding function type is unsupported. + ImportError: If required dependencies are not installed. + """ + try: + from chromadb.utils import embedding_functions + except ImportError as e: + raise ImportError( + "ChromaDB embedding functions not available. Ensure chromadb is properly installed." + ) from e + + config = self._config.embedding_function_config + + if isinstance(config, DefaultEmbeddingFunctionConfig): + return embedding_functions.DefaultEmbeddingFunction() + + elif isinstance(config, SentenceTransformerEmbeddingFunctionConfig): + try: + return embedding_functions.SentenceTransformerEmbeddingFunction(model_name=config.model_name) + except Exception as e: + raise ImportError( + f"Failed to create SentenceTransformer embedding function with model '{config.model_name}'. " + f"Ensure sentence-transformers is installed and the model is available. Error: {e}" + ) from e + + elif isinstance(config, OpenAIEmbeddingFunctionConfig): + try: + return embedding_functions.OpenAIEmbeddingFunction(api_key=config.api_key, model_name=config.model_name) + except Exception as e: + raise ImportError( + f"Failed to create OpenAI embedding function with model '{config.model_name}'. " + f"Ensure openai is installed and API key is valid. Error: {e}" + ) from e + + elif isinstance(config, CustomEmbeddingFunctionConfig): + try: + return config.function(**config.params) + except Exception as e: + raise ValueError(f"Failed to create custom embedding function. Error: {e}") from e + + else: + raise ValueError(f"Unsupported embedding function config type: {type(config)}") + def _ensure_initialized(self) -> None: """Ensure ChromaDB client and collection are initialized.""" if self._client is None: @@ -172,8 +268,14 @@ def _ensure_initialized(self) -> None: if self._collection is None: try: + # Create embedding function + embedding_function = self._create_embedding_function() + + # Create or get collection with embedding function self._collection = self._client.get_or_create_collection( - name=self._config.collection_name, metadata={"distance_metric": self._config.distance_metric} + name=self._config.collection_name, + metadata={"distance_metric": self._config.distance_metric}, + embedding_function=embedding_function, ) except Exception as e: logger.error(f"Failed to get/create collection: {e}") @@ -209,67 +311,6 @@ async def update_context( self, model_context: ChatCompletionContext, ) -> UpdateContextResult: - """ - Update the model context with relevant memory content. - - This method retrieves memory content relevant to the last message in the context - and adds it as a system message. It serves as the primary customization point for - how retrieved memories are incorporated into the conversation. - - By default, this implementation: - 1. Uses the last message as a query to find semantically similar memories - 2. Formats retrieved memories as a numbered list - 3. Adds them to the context as a system message - - For custom memory formatting, extend this class and override this method. - - Args: - model_context (ChatCompletionContext): The model context to update with relevant memories. - - Returns: - UpdateContextResult: Object containing the memories that were used to update the context. - - Example: - - .. code-block:: python - - from autogen_core.memory import Memory, MemoryContent, MemoryQueryResult, UpdateContextResult - from autogen_core.model_context import ChatCompletionContext - from autogen_core.models import SystemMessage - from autogen_ext.memory.chromadb import ChromaDBVectorMemory, PersistentChromaDBVectorMemoryConfig - - - class CustomVectorMemory(ChromaDBVectorMemory): - async def update_context( - self, - model_context: ChatCompletionContext, - ) -> UpdateContextResult: - # Get the last message to use as query - messages = await model_context.get_messages() - if not messages: - return UpdateContextResult(memories=MemoryQueryResult(results=[])) - - # Get query results - last_message = messages[-1] - query_text = last_message.content if isinstance(last_message.content, str) else str(last_message) - query_results = await self.query(query_text) - - if query_results.results: - # Custom formatting based on memory category - memory_strings = [] - for memory in query_results.results: - category = memory.metadata.get("category", "general") - if category == "preferences": - memory_strings.append(f"User Preference: {memory.content}") - else: - memory_strings.append(f"Memory: {memory.content}") - - # Add to context with custom header - memory_context = "IMPORTANT USER INFORMATION:\n" + "\n".join(memory_strings) - await model_context.add_message(SystemMessage(content=memory_context)) - - return UpdateContextResult(memories=query_results) - """ messages = await model_context.get_messages() if not messages: return UpdateContextResult(memories=MemoryQueryResult(results=[])) @@ -292,7 +333,6 @@ async def update_context( return UpdateContextResult(memories=query_results) async def add(self, content: MemoryContent, cancellation_token: CancellationToken | None = None) -> None: - """Add a memory content to ChromaDB.""" self._ensure_initialized() if self._collection is None: raise RuntimeError("Failed to initialize ChromaDB") @@ -318,7 +358,6 @@ async def query( cancellation_token: CancellationToken | None = None, **kwargs: Any, ) -> MemoryQueryResult: - """Query memory content based on vector similarity.""" self._ensure_initialized() if self._collection is None: raise RuntimeError("Failed to initialize ChromaDB") @@ -331,7 +370,7 @@ async def query( results = self._collection.query( query_texts=[query_text], n_results=self._config.k, - include=[IncludeEnum.documents, IncludeEnum.metadatas, IncludeEnum.distances], + include=["documents", "metadatas", "distances"], **kwargs, ) @@ -378,7 +417,6 @@ async def query( raise async def clear(self) -> None: - """Clear all entries from memory.""" self._ensure_initialized() if self._collection is None: raise RuntimeError("Failed to initialize ChromaDB") @@ -397,7 +435,6 @@ async def close(self) -> None: self._client = None async def reset(self) -> None: - """Reset the memory by deleting all data.""" self._ensure_initialized() if not self._config.allow_reset: raise RuntimeError("Reset not allowed. Set allow_reset=True in config to enable.") diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/mem0/__init__.py b/python/packages/autogen-ext/src/autogen_ext/memory/mem0/__init__.py new file mode 100644 index 000000000000..2f1af25679c0 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/memory/mem0/__init__.py @@ -0,0 +1,6 @@ +from ._mem0 import Mem0Memory, Mem0MemoryConfig + +__all__ = [ + "Mem0Memory", + "Mem0MemoryConfig", +] diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/mem0/_mem0.py b/python/packages/autogen-ext/src/autogen_ext/memory/mem0/_mem0.py new file mode 100644 index 000000000000..48dfaf0b109f --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/memory/mem0/_mem0.py @@ -0,0 +1,449 @@ +import io +import logging +import uuid +from contextlib import redirect_stderr, redirect_stdout +from datetime import datetime +from typing import Any, Dict, List, Optional, TypedDict, cast + +from autogen_core import CancellationToken, Component, ComponentBase +from autogen_core.memory import Memory, MemoryContent, MemoryQueryResult, UpdateContextResult +from autogen_core.model_context import ChatCompletionContext +from autogen_core.models import SystemMessage +from mem0 import Memory as Memory0 +from mem0 import MemoryClient +from pydantic import BaseModel, Field +from typing_extensions import Self + +logger = logging.getLogger(__name__) +logging.getLogger("chromadb").setLevel(logging.ERROR) + + +class Mem0MemoryConfig(BaseModel): + """Configuration for Mem0Memory component.""" + + user_id: Optional[str] = Field( + default=None, description="User ID for memory operations. If not provided, a UUID will be generated." + ) + limit: int = Field(default=10, description="Maximum number of results to return in memory queries.") + is_cloud: bool = Field(default=True, description="Whether to use cloud Mem0 client (True) or local client (False).") + api_key: Optional[str] = Field( + default=None, description="API key for cloud Mem0 client. Required if is_cloud=True." + ) + config: Optional[Dict[str, Any]] = Field( + default=None, description="Configuration dictionary for local Mem0 client. Required if is_cloud=False." + ) + + +class MemoryResult(TypedDict, total=False): + memory: str + score: float + metadata: Dict[str, Any] + created_at: str + updated_at: str + categories: List[str] + + +# pyright: reportGeneralTypeIssues=false +class Mem0Memory(Memory, Component[Mem0MemoryConfig], ComponentBase[Mem0MemoryConfig]): + """Mem0 memory implementation for AutoGen. + + This component integrates with Mem0.ai's memory system, providing an implementation + of AutoGen's Memory interface. It supports both cloud and local backends through the + mem0ai Python package. + + To use this component, you need to have the `mem0` (for cloud-only) or `mem0-local` (for local) + extra installed for the `autogen-ext` package: + + .. code-block:: bash + + pip install -U "autogen-ext[mem0]" # For cloud-based Mem0 + pip install -U "autogen-ext[mem0-local]" # For local Mem0 + + The memory component can store and retrieve information that agents need to remember + across conversations. It also provides context updating for language models with + relevant memories. + + Examples: + + .. code-block:: python + + import asyncio + from autogen_ext.memory.mem0 import Mem0Memory + from autogen_core.memory import MemoryContent + + + async def main() -> None: + # Create a local Mem0Memory (no API key required) + memory = Mem0Memory( + is_cloud=False, + config={"path": ":memory:"}, # Use in-memory storage for testing + ) + print("Memory initialized successfully!") + + # Add something to memory + test_content = "User likes the color blue." + await memory.add(MemoryContent(content=test_content, mime_type="text/plain")) + print(f"Added content: {test_content}") + + # Retrieve memories with a search query + results = await memory.query("What color does the user like?") + print(f"Query results: {len(results.results)} found") + + for i, result in enumerate(results.results): + print(f"Result {i+1}: {result}") + + + asyncio.run(main()) + + Output: + + .. code-block:: text + + Memory initialized successfully! + Added content: User likes the color blue. + Query results: 1 found + Result 1: content='User likes the color blue' mime_type='text/plain' metadata={'score': 0.6977155806281953, 'created_at': datetime.datetime(2025, 7, 6, 17, 25, 18, 754725, tzinfo=datetime.timezone(datetime.timedelta(days=-1, seconds=61200)))} + + Using it with an :class:`~autogen_agentchat.agents.AssistantAgent`: + + .. code-block:: python + + import asyncio + from autogen_agentchat.agents import AssistantAgent + from autogen_core.memory import MemoryContent + from autogen_ext.memory.mem0 import Mem0Memory + from autogen_ext.models.openai import OpenAIChatCompletionClient + + + async def main() -> None: + # Create a model client + model_client = OpenAIChatCompletionClient(model="gpt-4.1") + + # Create a Mem0 memory instance + memory = Mem0Memory( + user_id="user123", + is_cloud=False, + config={"path": ":memory:"}, # Use in-memory storage for testing + ) + + # Add something to memory + test_content = "User likes the color blue." + await memory.add(MemoryContent(content=test_content, mime_type="text/plain")) + + # Create an assistant agent with Mem0 memory + agent = AssistantAgent( + name="assistant", + model_client=model_client, + memory=[memory], + system_message="You are a helpful assistant that remembers user preferences.", + ) + + # Run a sample task + result = await agent.run(task="What color does the user like?") + print(result.messages[-1].content) # type: ignore + + + asyncio.run(main()) + + Output: + + .. code-block:: text + + User likes the color blue. + + Args: + user_id: Optional user ID for memory operations. If not provided, a UUID will be generated. + limit: Maximum number of results to return in memory queries. + is_cloud: Whether to use cloud Mem0 client (True) or local client (False). + api_key: API key for cloud Mem0 client. It will read from the environment MEM0_API_KEY if not provided. + config: Configuration dictionary for local Mem0 client. Required if is_cloud=False. + """ + + component_type = "memory" + component_provider_override = "autogen_ext.memory.mem0.Mem0Memory" + component_config_schema = Mem0MemoryConfig + + def __init__( + self, + user_id: Optional[str] = None, + limit: int = 10, + is_cloud: bool = True, + api_key: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, + ) -> None: + # Validate parameters + if not is_cloud and config is None: + raise ValueError("config is required when using local Mem0 client (is_cloud=False)") + + # Initialize instance variables + self._user_id = user_id or str(uuid.uuid4()) + self._limit = limit + self._is_cloud = is_cloud + self._api_key = api_key + self._config = config + + # Initialize client + if self._is_cloud: + self._client = MemoryClient(api_key=self._api_key) + else: + assert self._config is not None + config_dict = self._config + self._client = Memory0.from_config(config_dict=config_dict) # type: ignore + + @property + def user_id(self) -> str: + """Get the user ID for memory operations.""" + return self._user_id + + @property + def limit(self) -> int: + """Get the maximum number of results to return in memory queries.""" + return self._limit + + @property + def is_cloud(self) -> bool: + """Check if the Mem0 client is cloud-based.""" + return self._is_cloud + + @property + def config(self) -> Optional[Dict[str, Any]]: + """Get the configuration for the Mem0 client.""" + return self._config + + async def add( + self, + content: MemoryContent, + cancellation_token: Optional[CancellationToken] = None, + ) -> None: + """Add content to memory. + + Args: + content: The memory content to add. + cancellation_token: Optional token to cancel operation. + + Raises: + Exception: If there's an error adding content to mem0 memory. + """ + # Extract content based on mime type + if hasattr(content, "content") and hasattr(content, "mime_type"): + if content.mime_type in ["text/plain", "text/markdown"]: + message = str(content.content) + elif content.mime_type == "application/json": + # Convert JSON content to string representation + if isinstance(content.content, str): + message = content.content + else: + # Convert dict or other JSON serializable objects to string + import json + + message = json.dumps(content.content) + else: + message = str(content.content) + + # Extract metadata + metadata = content.metadata or {} + else: + # Handle case where content is directly provided as string + message = str(content) + metadata = {} + + # Check if operation is cancelled + if cancellation_token is not None and cancellation_token.cancelled: # type: ignore + return + + # Add to mem0 client + try: + user_id = metadata.pop("user_id", self._user_id) + # Suppress warning messages from mem0 MemoryClient + kwargs = {} if self._client.__class__.__name__ == "Memory" else {"output_format": "v1.1"} + with redirect_stdout(io.StringIO()), redirect_stderr(io.StringIO()): + self._client.add([{"role": "user", "content": message}], user_id=user_id, metadata=metadata, **kwargs) # type: ignore + except Exception as e: + # Log the error but don't crash + logger.error(f"Error adding to mem0 memory: {str(e)}") + raise + + async def query( + self, + query: str | MemoryContent = "", + cancellation_token: Optional[CancellationToken] = None, + **kwargs: Any, + ) -> MemoryQueryResult: + """Query memory for relevant content. + + Args: + query: The query to search for, either as string or MemoryContent. + cancellation_token: Optional token to cancel operation. + **kwargs: Additional query parameters to pass to mem0. + + Returns: + MemoryQueryResult containing search results. + """ + # Extract query text + if isinstance(query, str): + query_text = query + elif hasattr(query, "content"): + query_text = str(query.content) + else: + query_text = str(query) + + # Check if operation is cancelled + if ( + cancellation_token + and hasattr(cancellation_token, "cancelled") + and getattr(cancellation_token, "cancelled", False) + ): + return MemoryQueryResult(results=[]) + + try: + limit = kwargs.pop("limit", self._limit) + with redirect_stdout(io.StringIO()), redirect_stderr(io.StringIO()): + # Query mem0 client + results = self._client.search( # type: ignore + query_text, + user_id=self._user_id, + limit=limit, + **kwargs, + ) + + # Type-safe handling of results + if isinstance(results, dict) and "results" in results: + result_list = cast(List[MemoryResult], results["results"]) + else: + result_list = cast(List[MemoryResult], results) + + # Convert results to MemoryContent objects + memory_contents: List[MemoryContent] = [] + for result in result_list: + content_text = result.get("memory", "") + metadata: Dict[str, Any] = {} + + if "metadata" in result and result["metadata"]: + metadata = result["metadata"] + + # Add relevant fields to metadata + if "score" in result: + metadata["score"] = result["score"] + + # For created_at + if "created_at" in result and result.get("created_at"): + try: + metadata["created_at"] = datetime.fromisoformat(result["created_at"]) + except (ValueError, TypeError): + pass + + # For updated_at + if "updated_at" in result and result.get("updated_at"): + try: + metadata["updated_at"] = datetime.fromisoformat(result["updated_at"]) + except (ValueError, TypeError): + pass + + # For categories + if "categories" in result and result.get("categories"): + metadata["categories"] = result["categories"] + + # Create MemoryContent object + memory_content = MemoryContent( + content=content_text, + mime_type="text/plain", # Default to text/plain + metadata=metadata, + ) + memory_contents.append(memory_content) + + return MemoryQueryResult(results=memory_contents) + + except Exception as e: + # Log the error but return empty results + logger.error(f"Error querying mem0 memory: {str(e)}") + return MemoryQueryResult(results=[]) + + async def update_context( + self, + model_context: ChatCompletionContext, + ) -> UpdateContextResult: + """Update the model context with relevant memories. + + This method retrieves the conversation history from the model context, + uses the last message as a query to find relevant memories, and then + adds those memories to the context as a system message. + + Args: + model_context: The model context to update. + + Returns: + UpdateContextResult containing memories added to the context. + """ + # Get messages from context + messages = await model_context.get_messages() + if not messages: + return UpdateContextResult(memories=MemoryQueryResult(results=[])) + + # Use the last message as query + last_message = messages[-1] + query_text = last_message.content if isinstance(last_message.content, str) else str(last_message) + + # Query memory + query_results = await self.query(query_text, limit=self._limit) + + # If we have results, add them to the context + if query_results.results: + # Format memories as numbered list + memory_strings = [f"{i}. {str(memory.content)}" for i, memory in enumerate(query_results.results, 1)] + memory_context = "\nRelevant memories:\n" + "\n".join(memory_strings) + + # Add as system message + await model_context.add_message(SystemMessage(content=memory_context)) + + return UpdateContextResult(memories=query_results) + + async def clear(self) -> None: + """Clear all content from memory for the current user. + + Raises: + Exception: If there's an error clearing mem0 memory. + """ + try: + self._client.delete_all(user_id=self._user_id) # type: ignore + except Exception as e: + logger.error(f"Error clearing mem0 memory: {str(e)}") + raise + + async def close(self) -> None: + """Clean up resources if needed. + + This is a no-op for Mem0 clients as they don't require explicit cleanup. + """ + pass + + @classmethod + def _from_config(cls, config: Mem0MemoryConfig) -> Self: + """Create instance from configuration. + + Args: + config: Configuration for Mem0Memory component. + + Returns: + A new Mem0Memory instance. + """ + return cls( + user_id=config.user_id, + limit=config.limit, + is_cloud=config.is_cloud, + api_key=config.api_key, + config=config.config, + ) + + def _to_config(self) -> Mem0MemoryConfig: + """Convert instance to configuration. + + Returns: + Configuration representing this Mem0Memory instance. + """ + return Mem0MemoryConfig( + user_id=self._user_id, + limit=self._limit, + is_cloud=self._is_cloud, + api_key=self._api_key, + config=self._config, + ) diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/redis/__init__.py b/python/packages/autogen-ext/src/autogen_ext/memory/redis/__init__.py new file mode 100644 index 000000000000..606cf2d46178 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/memory/redis/__init__.py @@ -0,0 +1,9 @@ +from ._redis_memory import ( + RedisMemory, + RedisMemoryConfig, +) + +__all__ = [ + "RedisMemoryConfig", + "RedisMemory", +] diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/redis/_redis_memory.py b/python/packages/autogen-ext/src/autogen_ext/memory/redis/_redis_memory.py new file mode 100644 index 000000000000..98cb19e18837 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/memory/redis/_redis_memory.py @@ -0,0 +1,325 @@ +import logging +from typing import Any, List, Literal + +from autogen_core import CancellationToken, Component +from autogen_core.memory import Memory, MemoryContent, MemoryMimeType, MemoryQueryResult, UpdateContextResult +from autogen_core.model_context import ChatCompletionContext +from autogen_core.models import SystemMessage +from pydantic import BaseModel, Field + +logger = logging.getLogger(__name__) + +try: + from redis import Redis + from redisvl.extensions.message_history import SemanticMessageHistory + from redisvl.utils.utils import deserialize, serialize +except ImportError as e: + raise ImportError("To use Redis Memory RedisVL must be installed. Run `pip install autogen-ext[redisvl]`") from e + + +class RedisMemoryConfig(BaseModel): + """ + Configuration for Redis-based vector memory. + + This class defines the configuration options for using Redis as a vector memory store, + supporting semantic memory. It allows customization of the Redis connection, index settings, + similarity search parameters, and embedding model. + """ + + redis_url: str = Field(default="redis://localhost:6379", description="url of the Redis instance") + index_name: str = Field(default="chat_history", description="Name of the Redis collection") + prefix: str = Field(default="memory", description="prefix of the Redis collection") + distance_metric: Literal["cosine", "ip", "l2"] = "cosine" + algorithm: Literal["flat", "hnsw"] = "flat" + top_k: int = Field(default=10, description="Number of results to return in queries") + datatype: Literal["uint8", "int8", "float16", "float32", "float64", "bfloat16"] = "float32" + distance_threshold: float = Field(default=0.7, description="Minimum similarity score threshold") + model_name: str | None = Field( + default="sentence-transformers/all-mpnet-base-v2", description="Embedding model name" + ) + + +class RedisMemory(Memory, Component[RedisMemoryConfig]): + """ + Store and retrieve memory using vector similarity search powered by RedisVL. + + `RedisMemory` provides a vector-based memory implementation that uses RedisVL for storing and + retrieving content based on semantic similarity. It enhances agents with the ability to recall + contextually relevant information during conversations by leveraging vector embeddings to find + similar content. + + This implementation requires the RedisVL extra to be installed. Install with: + + .. code-block:: bash + + pip install "autogen-ext[redisvl]" + + Additionally, you will need access to a Redis instance. + To run a local instance of redis in docker: + + .. code-block:: bash + + docker run -d --name redis -p 6379:6379 redis:8 + + To download and run Redis locally: + + .. code-block:: bash + + curl -fsSL https://packages.redis.io/gpg | sudo gpg --dearmor -o /usr/share/keyrings/redis-archive-keyring.gpg + echo "deb [signed-by=/usr/share/keyrings/redis-archive-keyring.gpg] https://packages.redis.io/deb $(lsb_release -cs) main" | sudo tee /etc/apt/sources.list.d/redis.list + sudo apt-get update > /dev/null 2>&1 + sudo apt-get install redis-server > /dev/null 2>&1 + redis-server --daemonize yes + + Args: + config (RedisMemoryConfig | None): Configuration for the Redis memory. + If None, defaults to a RedisMemoryConfig with recommended settings. + + Example: + + .. code-block:: python + + from logging import WARNING, getLogger + + import asyncio + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.ui import Console + from autogen_core.memory import MemoryContent, MemoryMimeType + from autogen_ext.memory.redis import RedisMemory, RedisMemoryConfig + from autogen_ext.models.openai import OpenAIChatCompletionClient + + logger = getLogger() + logger.setLevel(WARNING) + + + # Define tool to use + async def get_weather(city: str, units: str = "imperial") -> str: + if units == "imperial": + return f"The weather in {city} is 73 °F and Sunny." + elif units == "metric": + return f"The weather in {city} is 23 °C and Sunny." + else: + return f"Sorry, I don't know the weather in {city}." + + + async def main(): + # Initailize Redis memory + redis_memory = RedisMemory( + config=RedisMemoryConfig( + redis_url="redis://localhost:6379", + index_name="chat_history", + prefix="memory", + ) + ) + + # Add user preferences to memory + await redis_memory.add( + MemoryContent( + content="The weather should be in metric units", + mime_type=MemoryMimeType.TEXT, + metadata={"category": "preferences", "type": "units"}, + ) + ) + + await redis_memory.add( + MemoryContent( + content="Meal recipe must be vegan", + mime_type=MemoryMimeType.TEXT, + metadata={"category": "preferences", "type": "dietary"}, + ) + ) + + model_client = OpenAIChatCompletionClient( + model="gpt-4o", + ) + + # Create assistant agent with ChromaDB memory + assistant_agent = AssistantAgent( + name="assistant_agent", + model_client=model_client, + tools=[get_weather], + memory=[redis_memory], + ) + + stream = assistant_agent.run_stream(task="What is the weather in New York?") + await Console(stream) + + await model_client.close() + await redis_memory.close() + + + asyncio.run(main()) + + Output: + + .. code-block:: text + + ---------- TextMessage (user) ---------- + What is the weather in New York? + ---------- MemoryQueryEvent (assistant_agent) ---------- + [MemoryContent(content='The weather should be in metric units', mime_type=, metadata={'category': 'preferences', 'type': 'units'})] + ---------- ToolCallRequestEvent (assistant_agent) ---------- + [FunctionCall(id='call_tyCPvPPAV4SHWhtfpM6UMemr', arguments='{"city":"New York","units":"metric"}', name='get_weather')] + ---------- ToolCallExecutionEvent (assistant_agent) ---------- + [FunctionExecutionResult(content='The weather in New York is 23 °C and Sunny.', name='get_weather', call_id='call_tyCPvPPAV4SHWhtfpM6UMemr', is_error=False)] + ---------- ToolCallSummaryMessage (assistant_agent) ---------- + The weather in New York is 23 °C and Sunny. + + """ + + component_config_schema = RedisMemoryConfig + component_provider_override = "autogen_ext.memory.redis_memory.RedisMemory" + + def __init__(self, config: RedisMemoryConfig | None = None) -> None: + """Initialize RedisMemory.""" + self.config = config or RedisMemoryConfig() + client = Redis.from_url(url=self.config.redis_url) # type: ignore[reportUknownMemberType] + + self.message_history = SemanticMessageHistory(name=self.config.index_name, redis_client=client) + + async def update_context( + self, + model_context: ChatCompletionContext, + ) -> UpdateContextResult: + """ + Update the model context with relevant memory content. + + This method retrieves memory content relevant to the last message in the context + and adds it as a system message. This implementation uses the last message in the context + as a query to find semantically similar memories and adds them all to the context as a + single system message. + + Args: + model_context (ChatCompletionContext): The model context to update with relevant + memories. + + Returns: + UpdateContextResult: Object containing the memories that were used to update the + context. + """ + messages = await model_context.get_messages() + if messages: + last_message = str(messages[-1].content) + else: + last_message = "" + + query_results = await self.query(last_message) + + stringified_messages = "\n\n".join([str(m.content) for m in query_results.results]) + + await model_context.add_message(SystemMessage(content=stringified_messages)) + + return UpdateContextResult(memories=query_results) + + async def add(self, content: MemoryContent, cancellation_token: CancellationToken | None = None) -> None: + """Add a memory content object to Redis. + + .. note:: + + To perform semantic search over stored memories RedisMemory creates a vector embedding + from the content field of a MemoryContent object. This content is assumed to be text, + JSON, or Markdown, and is passed to the vector embedding model specified in + RedisMemoryConfig. + + Args: + content (MemoryContent): The memory content to store within Redis. + cancellation_token (CancellationToken): Token passed to cease operation. Not used. + """ + if content.mime_type == MemoryMimeType.TEXT: + memory_content = content.content + mime_type = "text/plain" + elif content.mime_type == MemoryMimeType.JSON: + memory_content = serialize(content.content) + mime_type = "application/json" + elif content.mime_type == MemoryMimeType.MARKDOWN: + memory_content = content.content + mime_type = "text/markdown" + else: + raise NotImplementedError( + f"Error: {content.mime_type} is not supported. Only MemoryMimeType.TEXT, MemoryMimeType.JSON, and MemoryMimeType.MARKDOWN are currently supported." + ) + metadata = {"mime_type": mime_type} + metadata.update(content.metadata if content.metadata else {}) + self.message_history.add_message( + {"role": "user", "content": memory_content, "tool_call_id": serialize(metadata)} # type: ignore[reportArgumentType] + ) + + async def query( + self, + query: str | MemoryContent, + cancellation_token: CancellationToken | None = None, + **kwargs: Any, + ) -> MemoryQueryResult: + """Query memory content based on semantic vector similarity. + + .. note:: + + RedisMemory.query() supports additional keyword arguments to improve query performance. + top_k (int): The maximum number of relevant memories to include. Defaults to 10. + distance_threshold (float): The maximum distance in vector space to consider a memory + semantically similar when performining cosine similarity search. Defaults to 0.7. + + Args: + query (str | MemoryContent): query to perform vector similarity search with. If a + string is passed, a vector embedding is created from it with the model specified + in the RedisMemoryConfig. If a MemoryContent object is passed, the content field + of this object is extracted and a vector embedding is created from it with the + model specified in the RedisMemoryConfig. + cancellation_token (CancellationToken): Token passed to cease operation. Not used. + + Returns: + memoryQueryResult: Object containing memories relevant to the provided query. + """ + # get the query string, or raise an error for unsupported MemoryContent types + if isinstance(query, str): + prompt = query + elif isinstance(query, MemoryContent): + if query.mime_type in (MemoryMimeType.TEXT, MemoryMimeType.MARKDOWN): + prompt = str(query.content) + elif query.mime_type == MemoryMimeType.JSON: + prompt = serialize(query.content) + else: + raise NotImplementedError( + f"Error: {query.mime_type} is not supported. Only MemoryMimeType.TEXT, MemoryMimeType.JSON, MemoryMimeType.MARKDOWN are currently supported." + ) + else: + raise TypeError("'query' must be either a string or MemoryContent") + + top_k = kwargs.pop("top_k", self.config.top_k) + distance_threshold = kwargs.pop("distance_threshold", self.config.distance_threshold) + + results = self.message_history.get_relevant( + prompt=prompt, # type: ignore[reportArgumentType] + top_k=top_k, + distance_threshold=distance_threshold, + raw=False, + ) + + memories: List[MemoryContent] = [] + for result in results: + metadata = deserialize(result["tool_call_id"]) # type: ignore[reportArgumentType] + mime_type = MemoryMimeType(metadata.pop("mime_type")) + if mime_type in (MemoryMimeType.TEXT, MemoryMimeType.MARKDOWN): + memory_content = result["content"] # type: ignore[reportArgumentType] + elif mime_type == MemoryMimeType.JSON: + memory_content = deserialize(result["content"]) # type: ignore[reportArgumentType] + else: + raise NotImplementedError( + f"Error: {mime_type} is not supported. Only MemoryMimeType.TEXT, MemoryMimeType.JSON, and MemoryMimeType.MARKDOWN are currently supported." + ) + memory = MemoryContent( + content=memory_content, # type: ignore[reportArgumentType] + mime_type=mime_type, + metadata=metadata, + ) + memories.append(memory) # type: ignore[reportUknownMemberType] + + return MemoryQueryResult(results=memories) # type: ignore[reportUknownMemberType] + + async def clear(self) -> None: + """Clear all entries from memory, preserving the RedisMemory resources.""" + self.message_history.clear() + + async def close(self) -> None: + """Clears all entries from memory, and cleans up Redis client, index and resources.""" + self.message_history.delete() diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py index 75365d781eb3..f31e7b1c0b72 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py @@ -1,14 +1,25 @@ -from ._anthropic_client import AnthropicChatCompletionClient, BaseAnthropicChatCompletionClient +from ._anthropic_client import ( + AnthropicBedrockChatCompletionClient, + AnthropicChatCompletionClient, + BaseAnthropicChatCompletionClient, +) from .config import ( + AnthropicBedrockClientConfiguration, + AnthropicBedrockClientConfigurationConfigModel, AnthropicClientConfiguration, AnthropicClientConfigurationConfigModel, + BedrockInfo, CreateArgumentsConfigModel, ) __all__ = [ "AnthropicChatCompletionClient", + "AnthropicBedrockChatCompletionClient", "BaseAnthropicChatCompletionClient", "AnthropicClientConfiguration", + "AnthropicBedrockClientConfiguration", "AnthropicClientConfigurationConfigModel", + "AnthropicBedrockClientConfigurationConfigModel", "CreateArgumentsConfigModel", + "BedrockInfo", ] diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py index 5788f1bf44e3..3f78f86ded64 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py @@ -5,13 +5,12 @@ import logging import re import warnings - -# from asyncio import Task from typing import ( Any, AsyncGenerator, Coroutine, Dict, + Iterable, List, Literal, Mapping, @@ -20,10 +19,11 @@ Set, Union, cast, + overload, ) import tiktoken -from anthropic import AsyncAnthropic, AsyncStream +from anthropic import AsyncAnthropic, AsyncAnthropicBedrock, AsyncStream from anthropic.types import ( Base64ImageSourceParam, ContentBlock, @@ -61,11 +61,18 @@ validate_model_info, ) from autogen_core.tools import Tool, ToolSchema +from autogen_core.utils import extract_json_from_str from pydantic import BaseModel, SecretStr from typing_extensions import Self, Unpack from . import _model_info -from .config import AnthropicClientConfiguration, AnthropicClientConfigurationConfigModel +from .config import ( + AnthropicBedrockClientConfiguration, + AnthropicBedrockClientConfigurationConfigModel, + AnthropicClientConfiguration, + AnthropicClientConfigurationConfigModel, + BedrockInfo, +) logger = logging.getLogger(EVENT_LOGGER_NAME) trace_logger = logging.getLogger(TRACE_LOGGER_NAME) @@ -142,20 +149,66 @@ def get_mime_type_from_image(image: Image) -> Literal["image/jpeg", "image/png", return "image/jpeg" +def convert_tool_choice_anthropic(tool_choice: Tool | Literal["auto", "required", "none"]) -> Any: + """Convert tool_choice parameter to Anthropic API format. + + Args: + tool_choice: A single Tool object to force the model to use, "auto" to let the model choose any available tool, "required" to force tool usage, or "none" to disable tool usage. + + Returns: + Anthropic API compatible tool_choice value. + """ + if tool_choice == "none": + return {"type": "none"} + + if tool_choice == "auto": + return {"type": "auto"} + + if tool_choice == "required": + return {"type": "any"} # Anthropic uses "any" for required + + # Must be a Tool object + if isinstance(tool_choice, Tool): + return {"type": "tool", "name": tool_choice.schema["name"]} + else: + raise ValueError(f"tool_choice must be a Tool object, 'auto', 'required', or 'none', got {type(tool_choice)}") + + +@overload +def __empty_content_to_whitespace(content: str) -> str: ... + + +@overload +def __empty_content_to_whitespace(content: List[Any]) -> Iterable[Any]: ... + + +def __empty_content_to_whitespace( + content: Union[str, List[Union[str, Image]]], +) -> Union[str, Iterable[Any]]: + if isinstance(content, str) and not content.strip(): + return " " + elif isinstance(content, list) and not any(isinstance(x, str) and not x.strip() for x in content): + for idx, message in enumerate(content): + if isinstance(message, str) and not message.strip(): + content[idx] = " " + + return content + + def user_message_to_anthropic(message: UserMessage) -> MessageParam: assert_valid_name(message.source) if isinstance(message.content, str): return { "role": "user", - "content": message.content, + "content": __empty_content_to_whitespace(message.content), } else: blocks: List[Union[TextBlockParam, ImageBlockParam]] = [] for part in message.content: if isinstance(part, str): - blocks.append(TextBlockParam(type="text", text=part)) + blocks.append(TextBlockParam(type="text", text=__empty_content_to_whitespace(part))) elif isinstance(part, Image): blocks.append( ImageBlockParam( @@ -177,7 +230,7 @@ def user_message_to_anthropic(message: UserMessage) -> MessageParam: def system_message_to_anthropic(message: SystemMessage) -> str: - return message.content + return __empty_content_to_whitespace(message.content) def assistant_message_to_anthropic(message: AssistantMessage) -> MessageParam: @@ -190,9 +243,13 @@ def assistant_message_to_anthropic(message: AssistantMessage) -> MessageParam: for func_call in message.content: # Parse the arguments and convert to dict if it's a JSON string args = func_call.arguments + args = __empty_content_to_whitespace(args) if isinstance(args, str): try: - args_dict = json.loads(args) + json_objs = extract_json_from_str(args) + if len(json_objs) != 1: + raise ValueError(f"Expected a single JSON object, but found {len(json_objs)}") + args_dict = json_objs[0] except json.JSONDecodeError: args_dict = {"text": args} else: @@ -386,7 +443,7 @@ def _add_usage(usage1: RequestUsage, usage2: RequestUsage) -> RequestUsage: class BaseAnthropicChatCompletionClient(ChatCompletionClient): def __init__( self, - client: AsyncAnthropic, + client: Any, *, create_args: Dict[str, Any], model_info: Optional[ModelInfo] = None, @@ -408,11 +465,71 @@ def __init__( self._total_usage = RequestUsage(prompt_tokens=0, completion_tokens=0) self._actual_usage = RequestUsage(prompt_tokens=0, completion_tokens=0) + def _serialize_message(self, message: MessageParam) -> Dict[str, Any]: + """Convert an Anthropic MessageParam to a JSON-serializable format.""" + if isinstance(message, dict): + result: Dict[str, Any] = {} + for key, value in message.items(): + if key == "content" and isinstance(value, list): + serialized_blocks: List[Any] = [] + for block in value: # type: ignore + if isinstance(block, BaseModel): + serialized_blocks.append(block.model_dump()) + else: + serialized_blocks.append(block) + result[key] = serialized_blocks + else: + result[key] = value + return result + else: + return {"role": "unknown", "content": str(message)} + + def _merge_system_messages(self, messages: Sequence[LLMMessage]) -> Sequence[LLMMessage]: + """ + Merge continuous system messages into a single message. + """ + _messages: List[LLMMessage] = [] + system_message_content = "" + _first_system_message_idx = -1 + _last_system_message_idx = -1 + # Index of the first system message for adding the merged system message at the correct position + for idx, message in enumerate(messages): + if isinstance(message, SystemMessage): + if _first_system_message_idx == -1: + _first_system_message_idx = idx + elif _last_system_message_idx + 1 != idx: + # That case, system message is not continuous + # Merge system messages only contiues system messages + raise ValueError("Multiple and Not continuous system messages are not supported") + system_message_content += message.content + "\n" + _last_system_message_idx = idx + else: + _messages.append(message) + system_message_content = system_message_content.rstrip() + if system_message_content != "": + system_message = SystemMessage(content=system_message_content) + _messages.insert(_first_system_message_idx, system_message) + messages = _messages + + return messages + + def _rstrip_last_assistant_message(self, messages: Sequence[LLMMessage]) -> Sequence[LLMMessage]: + """ + Remove the last assistant message if it is empty. + """ + # When Claude models last message is AssistantMessage, It could not end with whitespace + if isinstance(messages[-1], AssistantMessage): + if isinstance(messages[-1].content, str): + messages[-1].content = messages[-1].content.rstrip() + + return messages + async def create( self, messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -442,9 +559,14 @@ async def create( system_message = None anthropic_messages: List[MessageParam] = [] + # Merge continuous system messages into a single message + messages = self._merge_system_messages(messages) + messages = self._rstrip_last_assistant_message(messages) + for message in messages: if isinstance(message, SystemMessage): if system_message is not None: + # if that case, system message is must only one raise ValueError("Multiple system messages are not supported") system_message = to_anthropic_type(message) else: @@ -486,6 +608,36 @@ async def create( # anthropic requires tools to be present even if there is any tool use request_args["tools"] = self._last_used_tools + # Process tool_choice parameter + if isinstance(tool_choice, Tool): + if len(tools) == 0 and not has_tool_results: + raise ValueError("tool_choice specified but no tools provided") + + # Validate that the tool exists in the provided tools + tool_names_available: List[str] = [] + if len(tools) > 0: + for tool in tools: + if isinstance(tool, Tool): + tool_names_available.append(tool.schema["name"]) + else: + tool_names_available.append(tool["name"]) + else: + # Use last used tools names if available + for tool_param in self._last_used_tools: + tool_names_available.append(tool_param["name"]) + + # tool_choice is a single Tool object + tool_name = tool_choice.schema["name"] + if tool_name not in tool_names_available: + raise ValueError(f"tool_choice references '{tool_name}' but it's not in the available tools") + + # Convert to Anthropic format and add to request_args only if tools are provided + # According to Anthropic API, tool_choice may only be specified while providing tools + if len(tools) > 0 or has_tool_results: + converted_tool_choice = convert_tool_choice_anthropic(tool_choice) + if converted_tool_choice is not None: + request_args["tool_choice"] = converted_tool_choice + # Optional parameters for param in ["top_p", "top_k", "stop_sequences", "metadata"]: if param in create_args: @@ -504,10 +656,11 @@ async def create( prompt_tokens=result.usage.input_tokens, completion_tokens=result.usage.output_tokens, ) + serializable_messages: List[Dict[str, Any]] = [self._serialize_message(msg) for msg in anthropic_messages] logger.info( LLMCallEvent( - messages=cast(List[Dict[str, Any]], anthropic_messages), + messages=serializable_messages, response=result.model_dump(), prompt_tokens=usage.prompt_tokens, completion_tokens=usage.completion_tokens, @@ -570,6 +723,7 @@ async def create_stream( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -604,9 +758,14 @@ async def create_stream( system_message = None anthropic_messages: List[MessageParam] = [] + # Merge continuous system messages into a single message + messages = self._merge_system_messages(messages) + messages = self._rstrip_last_assistant_message(messages) + for message in messages: if isinstance(message, SystemMessage): if system_message is not None: + # if that case, system message is must only one raise ValueError("Multiple system messages are not supported") system_message = to_anthropic_type(message) else: @@ -649,6 +808,36 @@ async def create_stream( elif has_tool_results: request_args["tools"] = self._last_used_tools + # Process tool_choice parameter + if isinstance(tool_choice, Tool): + if len(tools) == 0 and not has_tool_results: + raise ValueError("tool_choice specified but no tools provided") + + # Validate that the tool exists in the provided tools + tool_names_available: List[str] = [] + if len(tools) > 0: + for tool in tools: + if isinstance(tool, Tool): + tool_names_available.append(tool.schema["name"]) + else: + tool_names_available.append(tool["name"]) + else: + # Use last used tools names if available + for last_used_tool in self._last_used_tools: + tool_names_available.append(last_used_tool["name"]) + + # tool_choice is a single Tool object + tool_name = tool_choice.schema["name"] + if tool_name not in tool_names_available: + raise ValueError(f"tool_choice references '{tool_name}' but it's not in the available tools") + + # Convert to Anthropic format and add to request_args only if tools are provided + # According to Anthropic API, tool_choice may only be specified while providing tools + if len(tools) > 0 or has_tool_results: + converted_tool_choice = convert_tool_choice_anthropic(tool_choice) + if converted_tool_choice is not None: + request_args["tool_choice"] = converted_tool_choice + # Optional parameters for param in ["top_p", "top_k", "stop_sequences", "metadata"]: if param in create_args: @@ -672,6 +861,7 @@ async def create_stream( stop_reason: Optional[str] = None first_chunk = True + serialized_messages: List[Dict[str, Any]] = [self._serialize_message(msg) for msg in anthropic_messages] # Process the stream async for chunk in stream: @@ -680,7 +870,7 @@ async def create_stream( # Emit the start event. logger.info( LLMStreamStartEvent( - messages=cast(List[Dict[str, Any]], anthropic_messages), + messages=serialized_messages, ) ) # Handle different event types @@ -1014,3 +1204,138 @@ def _from_config(cls, config: AnthropicClientConfigurationConfigModel) -> Self: copied_config["api_key"] = config.api_key.get_secret_value() return cls(**copied_config) + + +class AnthropicBedrockChatCompletionClient( + BaseAnthropicChatCompletionClient, Component[AnthropicBedrockClientConfigurationConfigModel] +): + """ + Chat completion client for Anthropic's Claude models on AWS Bedrock. + + Args: + model (str): The Claude model to use (e.g., "claude-3-sonnet-20240229", "claude-3-opus-20240229") + api_key (str, optional): Anthropic API key. Required if not in environment variables. + base_url (str, optional): Override the default API endpoint. + max_tokens (int, optional): Maximum tokens in the response. Default is 4096. + temperature (float, optional): Controls randomness. Lower is more deterministic. Default is 1.0. + top_p (float, optional): Controls diversity via nucleus sampling. Default is 1.0. + top_k (int, optional): Controls diversity via top-k sampling. Default is -1 (disabled). + model_info (ModelInfo, optional): The capabilities of the model. Required if using a custom model. + bedrock_info (BedrockInfo, optional): The capabilities of the model in bedrock. Required if using a model from AWS bedrock. + + To use this client, you must install the Anthropic extension: + + .. code-block:: bash + + pip install "autogen-ext[anthropic]" + + Example: + + .. code-block:: python + + import asyncio + from autogen_ext.models.anthropic import AnthropicBedrockChatCompletionClient, BedrockInfo + from autogen_core.models import UserMessage, ModelInfo + + + async def main(): + anthropic_client = AnthropicBedrockChatCompletionClient( + model="anthropic.claude-3-5-sonnet-20240620-v1:0", + temperature=0.1, + model_info=ModelInfo( + vision=False, function_calling=True, json_output=False, family="unknown", structured_output=True + ), + bedrock_info=BedrockInfo( + aws_access_key="", + aws_secret_key="", + aws_session_token="", + aws_region="", + ), + ) + + result = await anthropic_client.create([UserMessage(content="What is the capital of France?", source="user")]) # type: ignore + print(result) + + + if __name__ == "__main__": + asyncio.run(main()) + """ + + component_type = "model" + component_config_schema = AnthropicBedrockClientConfigurationConfigModel + component_provider_override = "autogen_ext.models.anthropic.AnthropicBedrockChatCompletionClient" + + def __init__(self, **kwargs: Unpack[AnthropicBedrockClientConfiguration]): + if "model" not in kwargs: + raise ValueError("model is required for AnthropicBedrockChatCompletionClient") + + self._raw_config: Dict[str, Any] = dict(kwargs).copy() + copied_args = dict(kwargs).copy() + + model_info: Optional[ModelInfo] = None + if "model_info" in kwargs: + model_info = kwargs["model_info"] + del copied_args["model_info"] + + bedrock_info: Optional[BedrockInfo] = None + if "bedrock_info" in kwargs: + bedrock_info = kwargs["bedrock_info"] + + if bedrock_info is None: + raise ValueError("bedrock_info is required for AnthropicBedrockChatCompletionClient") + + # Handle bedrock_info + aws_region = bedrock_info["aws_region"] + aws_access_key: Optional[str] = None + aws_secret_key: Optional[str] = None + aws_session_token: Optional[str] = None + if all(key in bedrock_info for key in ("aws_access_key", "aws_secret_key", "aws_session_token")): + aws_access_key = bedrock_info["aws_access_key"] + aws_secret_key = bedrock_info["aws_secret_key"] + aws_session_token = bedrock_info["aws_session_token"] + + client = AsyncAnthropicBedrock( + aws_access_key=aws_access_key, + aws_secret_key=aws_secret_key, + aws_session_token=aws_session_token, + aws_region=aws_region, + ) + create_args = _create_args_from_config(copied_args) + + super().__init__( + client=client, + create_args=create_args, + model_info=model_info, + ) + + def __getstate__(self) -> Dict[str, Any]: + state = self.__dict__.copy() + state["_client"] = None + return state + + def __setstate__(self, state: Dict[str, Any]) -> None: + self.__dict__.update(state) + self._client = _anthropic_client_from_config(state["_raw_config"]) + + def _to_config(self) -> AnthropicBedrockClientConfigurationConfigModel: + copied_config = self._raw_config.copy() + return AnthropicBedrockClientConfigurationConfigModel(**copied_config) + + @classmethod + def _from_config(cls, config: AnthropicBedrockClientConfigurationConfigModel) -> Self: + copied_config = config.model_copy().model_dump(exclude_none=True) + + # Handle api_key as SecretStr + if "api_key" in copied_config and isinstance(config.api_key, SecretStr): + copied_config["api_key"] = config.api_key.get_secret_value() + + # Handle bedrock_info as SecretStr + if "bedrock_info" in copied_config and isinstance(config.bedrock_info, dict): + copied_config["bedrock_info"] = { + "aws_access_key": config.bedrock_info["aws_access_key"].get_secret_value(), + "aws_secret_key": config.bedrock_info["aws_secret_key"].get_secret_value(), + "aws_session_token": config.bedrock_info["aws_session_token"].get_secret_value(), + "aws_region": config.bedrock_info["aws_region"], + } + + return cls(**copied_config) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_model_info.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_model_info.py index 059af5097632..81dc4f79638e 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_model_info.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_model_info.py @@ -6,6 +6,42 @@ # For Anthropic's Claude models based on: # https://docs.anthropic.com/claude/docs/models-overview _MODEL_INFO: Dict[str, ModelInfo] = { + # Claude 4 Opus + "claude-opus-4-20250514": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.CLAUDE_4_OPUS, + "structured_output": False, + "multiple_system_messages": False, + }, + # Claude 4 Opus latest alias + "claude-opus-4-0": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.CLAUDE_4_OPUS, + "structured_output": False, + "multiple_system_messages": False, + }, + # Claude 4 Sonnet + "claude-sonnet-4-20250514": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.CLAUDE_4_SONNET, + "structured_output": False, + "multiple_system_messages": False, + }, + # Claude 4 Sonnet latest alias + "claude-sonnet-4-0": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.CLAUDE_4_SONNET, + "structured_output": False, + "multiple_system_messages": False, + }, # Claude 3.7 Sonnet "claude-3-7-sonnet-20250219": { "vision": True, @@ -13,6 +49,7 @@ "json_output": True, "family": ModelFamily.CLAUDE_3_7_SONNET, "structured_output": False, + "multiple_system_messages": False, }, # Claude 3.7 Sonnet latest alias "claude-3-7-sonnet-latest": { @@ -21,6 +58,7 @@ "json_output": True, "family": ModelFamily.CLAUDE_3_7_SONNET, "structured_output": False, + "multiple_system_messages": False, }, # Claude 3 Opus (most powerful) "claude-3-opus-20240229": { @@ -29,6 +67,7 @@ "json_output": True, "family": ModelFamily.CLAUDE_3_5_SONNET, "structured_output": False, + "multiple_system_messages": False, }, # Claude 3 Sonnet (balanced) "claude-3-sonnet-20240229": { @@ -37,6 +76,7 @@ "json_output": True, "family": ModelFamily.CLAUDE_3_5_SONNET, "structured_output": False, + "multiple_system_messages": False, }, # Claude 3 Haiku (fastest) "claude-3-haiku-20240307": { @@ -45,6 +85,7 @@ "json_output": True, "family": ModelFamily.CLAUDE_3_5_SONNET, "structured_output": False, + "multiple_system_messages": False, }, # Claude 3.5 Sonnet "claude-3-5-sonnet-20240620": { @@ -53,6 +94,7 @@ "json_output": True, "family": ModelFamily.CLAUDE_3_5_SONNET, "structured_output": False, + "multiple_system_messages": False, }, # Claude Instant v1 (legacy) "claude-instant-1.2": { @@ -61,6 +103,7 @@ "json_output": True, "family": ModelFamily.CLAUDE_3_5_SONNET, "structured_output": False, + "multiple_system_messages": False, }, # Claude 2 (legacy) "claude-2.0": { @@ -69,6 +112,7 @@ "json_output": True, "family": ModelFamily.CLAUDE_3_5_SONNET, "structured_output": False, + "multiple_system_messages": False, }, # Claude 2.1 (legacy) "claude-2.1": { @@ -77,6 +121,7 @@ "json_output": True, "family": ModelFamily.CLAUDE_3_5_SONNET, "structured_output": False, + "multiple_system_messages": False, }, } diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py index 0324f66b13ff..22dfc067f318 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py @@ -2,7 +2,7 @@ from autogen_core.models import ModelCapabilities, ModelInfo # type: ignore from pydantic import BaseModel, SecretStr -from typing_extensions import TypedDict +from typing_extensions import Required, TypedDict class ResponseFormat(TypedDict): @@ -20,6 +20,22 @@ class CreateArguments(TypedDict, total=False): metadata: Optional[Dict[str, str]] +class BedrockInfo(TypedDict): + """BedrockInfo is a dictionary that contains information about a bedrock's properties. + It is expected to be used in the bedrock_info property of a model client. + + """ + + aws_access_key: Required[str] + """Access key for the aws account to gain bedrock model access""" + aws_secret_key: Required[str] + """Access secret key for the aws account to gain bedrock model access""" + aws_session_token: Required[str] + """aws session token for the aws account to gain bedrock model access""" + aws_region: Required[str] + """aws region for the aws account to gain bedrock model access""" + + class BaseAnthropicClientConfiguration(CreateArguments, total=False): api_key: str base_url: Optional[str] @@ -36,6 +52,10 @@ class AnthropicClientConfiguration(BaseAnthropicClientConfiguration, total=False tool_choice: Optional[Union[Literal["auto", "any", "none"], Dict[str, Any]]] +class AnthropicBedrockClientConfiguration(AnthropicClientConfiguration, total=False): + bedrock_info: BedrockInfo + + # Pydantic equivalents of the above TypedDicts class CreateArgumentsConfigModel(BaseModel): model: str @@ -61,3 +81,17 @@ class BaseAnthropicClientConfigurationConfigModel(CreateArgumentsConfigModel): class AnthropicClientConfigurationConfigModel(BaseAnthropicClientConfigurationConfigModel): tools: List[Dict[str, Any]] | None = None tool_choice: Union[Literal["auto", "any", "none"], Dict[str, Any]] | None = None + + +class BedrockInfoConfigModel(TypedDict): + aws_access_key: Required[SecretStr] + """Access key for the aws account to gain bedrock model access""" + aws_session_token: Required[SecretStr] + """aws session token for the aws account to gain bedrock model access""" + aws_region: Required[str] + """aws region for the aws account to gain bedrock model access""" + aws_secret_key: Required[SecretStr] + + +class AnthropicBedrockClientConfigurationConfigModel(AnthropicClientConfigurationConfigModel): + bedrock_info: BedrockInfoConfigModel | None = None diff --git a/python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py index ebf25404b056..a783d4965e17 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py @@ -3,7 +3,7 @@ import re from asyncio import Task from inspect import getfullargspec -from typing import Any, Dict, List, Mapping, Optional, Sequence, cast +from typing import Any, Dict, List, Literal, Mapping, Optional, Sequence, Union, cast from autogen_core import EVENT_LOGGER_NAME, CancellationToken, FunctionCall, Image from autogen_core.logging import LLMCallEvent, LLMStreamEndEvent, LLMStreamStartEvent @@ -28,6 +28,8 @@ ) from azure.ai.inference.models import ( ChatCompletions, + ChatCompletionsNamedToolChoice, + ChatCompletionsNamedToolChoiceFunction, ChatCompletionsToolCall, ChatCompletionsToolDefinition, CompletionsFinishReason, @@ -53,7 +55,7 @@ UserMessage as AzureUserMessage, ) from pydantic import BaseModel -from typing_extensions import AsyncGenerator, Union, Unpack +from typing_extensions import AsyncGenerator, Unpack from autogen_ext.models.azure.config import ( GITHUB_MODELS_ENDPOINT, @@ -218,7 +220,7 @@ class AzureAIChatCompletionClient(ChatCompletionClient): async def main(): client = AzureAIChatCompletionClient( model="Phi-4", - endpoint="https://models.inference.ai.azure.com", + endpoint="https://models.github.ai/inference", # To authenticate with the model you will need to generate a personal access token (PAT) in your GitHub settings. # Create your PAT token by following instructions here: https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens credential=AzureKeyCredential(os.environ["GITHUB_TOKEN"]), @@ -256,7 +258,7 @@ async def main(): async def main(): client = AzureAIChatCompletionClient( model="Phi-4", - endpoint="https://models.inference.ai.azure.com", + endpoint="https://models.github.ai/inference", # To authenticate with the model you will need to generate a personal access token (PAT) in your GitHub settings. # Create your PAT token by following instructions here: https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens credential=AzureKeyCredential(os.environ["GITHUB_TOKEN"]), @@ -309,7 +311,10 @@ def _validate_config(config: Dict[str, Any]) -> AzureAIChatCompletionClientConfi @staticmethod def _create_client(config: AzureAIChatCompletionClientConfig) -> ChatCompletionsClient: - return ChatCompletionsClient(**config) + # Only pass the parameters that ChatCompletionsClient accepts + # Remove 'model_info' and other client-specific parameters + client_config = {k: v for k, v in config.items() if k not in ("model_info",)} + return ChatCompletionsClient(**client_config) # type: ignore @staticmethod def _prepare_create_args(config: Mapping[str, Any]) -> Dict[str, Any]: @@ -356,6 +361,7 @@ async def create( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -376,6 +382,12 @@ async def create( task: Task[ChatCompletions] if len(tools) > 0: + if isinstance(tool_choice, Tool): + create_args["tool_choice"] = ChatCompletionsNamedToolChoice( + function=ChatCompletionsNamedToolChoiceFunction(name=tool_choice.name) + ) + else: + create_args["tool_choice"] = tool_choice converted_tools = convert_tools(tools) task = asyncio.create_task( # type: ignore self._client.complete(messages=azure_messages, tools=converted_tools, **create_args) # type: ignore @@ -451,6 +463,7 @@ async def create_stream( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -469,6 +482,12 @@ async def create_stream( azure_messages = [item for sublist in azure_messages_nested for item in sublist] if len(tools) > 0: + if isinstance(tool_choice, Tool): + create_args["tool_choice"] = ChatCompletionsNamedToolChoice( + function=ChatCompletionsNamedToolChoiceFunction(name=tool_choice.name) + ) + else: + create_args["tool_choice"] = tool_choice converted_tools = convert_tools(tools) task = asyncio.create_task( self._client.complete(messages=azure_messages, tools=converted_tools, stream=True, **create_args) @@ -603,11 +622,3 @@ def model_info(self) -> ModelInfo: @property def capabilities(self) -> ModelInfo: return self.model_info - - def __del__(self) -> None: - # TODO: This is a hack to close the open client - if hasattr(self, "_client"): - try: - asyncio.get_running_loop().create_task(self._client.close()) - except RuntimeError: - asyncio.run(self._client.close()) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/azure/config/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/azure/config/__init__.py index c90afe89702c..38cf34b5378a 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/azure/config/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/azure/config/__init__.py @@ -9,7 +9,7 @@ from azure.core.credentials import AzureKeyCredential from azure.core.credentials_async import AsyncTokenCredential -GITHUB_MODELS_ENDPOINT = "https://models.inference.ai.azure.com" +GITHUB_MODELS_ENDPOINT = "https://models.github.ai/inference" class JsonSchemaFormat(TypedDict, total=False): diff --git a/python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py b/python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py index dd34809b0281..07e3d89e83ea 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py @@ -1,7 +1,7 @@ import hashlib import json import warnings -from typing import Any, AsyncGenerator, List, Mapping, Optional, Sequence, Union, cast +from typing import Any, AsyncGenerator, List, Literal, Mapping, Optional, Sequence, Union, cast from autogen_core import CacheStore, CancellationToken, Component, ComponentModel, InMemoryStore from autogen_core.models import ( @@ -137,6 +137,7 @@ async def create( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -158,6 +159,7 @@ async def create( messages, tools=tools, json_output=json_output, + tool_choice=tool_choice, extra_create_args=extra_create_args, cancellation_token=cancellation_token, ) @@ -169,6 +171,7 @@ def create_stream( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -200,6 +203,7 @@ async def _generator() -> AsyncGenerator[Union[str, CreateResult], None]: messages, tools=tools, json_output=json_output, + tool_choice=tool_choice, extra_create_args=extra_create_args, cancellation_token=cancellation_token, ) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/llama_cpp/_llama_cpp_completion_client.py b/python/packages/autogen-ext/src/autogen_ext/models/llama_cpp/_llama_cpp_completion_client.py index 32df1b2f67ef..36b115668b49 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/llama_cpp/_llama_cpp_completion_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/llama_cpp/_llama_cpp_completion_client.py @@ -1,6 +1,7 @@ import asyncio import logging # added import import re +import warnings from typing import Any, AsyncGenerator, Dict, List, Literal, Mapping, Optional, Sequence, TypedDict, Union, cast from autogen_core import EVENT_LOGGER_NAME, CancellationToken, FunctionCall, MessageHandlerContext @@ -264,6 +265,7 @@ async def create( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", # None means do not override the default # A value means to override the client default - often specified in the constructor json_output: Optional[bool | type[BaseModel]] = None, @@ -302,6 +304,15 @@ async def create( elif json_output is not False and json_output is not None: raise ValueError("json_output must be a boolean, a BaseModel subclass or None.") + # Handle tool_choice parameter + if tool_choice != "auto": + warnings.warn( + "tool_choice parameter is specified but LlamaCppChatCompletionClient does not support it. " + "This parameter will be ignored.", + UserWarning, + stacklevel=2, + ) + if self.model_info["function_calling"]: # Run this in on the event loop to avoid blocking. response_future = asyncio.get_event_loop().run_in_executor( @@ -397,12 +408,21 @@ async def create_stream( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", # None means do not override the default # A value means to override the client default - often specified in the constructor json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, ) -> AsyncGenerator[Union[str, CreateResult], None]: + # Validate tool_choice parameter even though streaming is not implemented + if tool_choice != "auto" and tool_choice != "none": + if not self.model_info["function_calling"]: + raise ValueError("tool_choice specified but model does not support function calling") + if len(tools) == 0: + raise ValueError("tool_choice specified but no tools provided") + logger.warning("tool_choice parameter specified but may not be supported by llama-cpp-python") + raise NotImplementedError("Stream not yet implemented for LlamaCppChatCompletionClient") yield "" diff --git a/python/packages/autogen-ext/src/autogen_ext/models/ollama/_model_info.py b/python/packages/autogen-ext/src/autogen_ext/models/ollama/_model_info.py index 158a0008fc6c..dcc6596f5c5e 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/ollama/_model_info.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/ollama/_model_info.py @@ -141,7 +141,7 @@ }, "llama3.3": { "vision": False, - "function_calling": False, + "function_calling": True, "json_output": True, "family": ModelFamily.UNKNOWN, "structured_output": True, @@ -258,6 +258,20 @@ "family": ModelFamily.UNKNOWN, "structured_output": True, }, + "qwen2.5vl": { + "vision": True, + "function_calling": False, + "json_output": True, + "family": ModelFamily.UNKNOWN, + "structured_output": True, + }, + "qwen3": { + "vision": False, + "function_calling": True, + "json_output": True, + "family": ModelFamily.UNKNOWN, + "structured_output": True, + }, "snowflake-arctic-embed": { "vision": False, "function_calling": False, @@ -351,6 +365,20 @@ "qwen2.5-coder:0.5b": 32768, "qwen2.5-coder:1.5b": 32768, "qwen2.5-coder:3b": 32768, + "qwen2.5vl": 128000, + "qwen2.5vl:3b": 128000, + "qwen2.5vl:7b": 128000, + "qwen2.5vl:32b": 128000, + "qwen2.5vl:72b": 128000, + "qwen3": 40960, + "qwen3:0.6b": 40960, + "qwen3:1.7b": 40960, + "qwen3:4b": 40960, + "qwen3:8b": 40960, + "qwen3:14b": 40960, + "qwen3:30b": 40960, + "qwen3:32b": 40960, + "qwen3:235b": 40960, "snowflake-arctic-embed": 512, "starcoder2": 16384, "tinyllama": 2048, diff --git a/python/packages/autogen-ext/src/autogen_ext/models/ollama/_ollama_client.py b/python/packages/autogen-ext/src/autogen_ext/models/ollama/_ollama_client.py index 5bf9a263c376..fbd4300c10b0 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/ollama/_ollama_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/ollama/_ollama_client.py @@ -73,11 +73,30 @@ def _ollama_client_from_config(config: Mapping[str, Any]) -> AsyncClient: # Take a copy copied_config = dict(config).copy() - # Shave down the config to just the AzureOpenAIChatCompletionClient kwargs + # Shave down the config to just the AsyncClient kwargs ollama_config = {k: v for k, v in copied_config.items() if k in ollama_init_kwargs} return AsyncClient(**ollama_config) +LLM_CONTROL_PARAMS = { + "temperature", + "top_p", + "top_k", + "repeat_penalty", + "frequency_penalty", + "presence_penalty", + "mirostat", + "mirostat_eta", + "mirostat_tau", + "seed", + "num_ctx", + "num_predict", + "num_gpu", + "stop", + "tfs_z", + "typical_p", +} + ollama_chat_request_fields: dict[str, Any] = [m for m in inspect.getmembers(ChatRequest) if m[0] == "model_fields"][0][ 1 ] @@ -95,18 +114,31 @@ def _create_args_from_config(config: Mapping[str, Any]) -> Dict[str, Any]: DeprecationWarning, stacklevel=2, ) - create_args = {k.lower(): v for k, v in config.items() if k.lower() in OLLAMA_VALID_CREATE_KWARGS_KEYS} - dropped_keys = [k for k in config.keys() if k.lower() not in OLLAMA_VALID_CREATE_KWARGS_KEYS] - trace_logger.info(f"Dropped the following unrecognized keys from create_args: {dropped_keys}") + + create_args: Dict[str, Any] = {} + options_dict: Dict[str, Any] = {} + + if "options" in config: + if isinstance(config["options"], Mapping): + options_map: Mapping[str, Any] = config["options"] + options_dict = dict(options_map) + else: + options_dict = {} + + for k, v in config.items(): + k_lower = k.lower() + if k_lower in OLLAMA_VALID_CREATE_KWARGS_KEYS: + create_args[k_lower] = v + elif k_lower in LLM_CONTROL_PARAMS: + options_dict[k_lower] = v + trace_logger.info(f"Moving LLM control parameter '{k}' to options dict") + else: + trace_logger.info(f"Dropped unrecognized key from create_args: {k}") + + if options_dict: + create_args["options"] = options_dict return create_args - # create_args = {k: v for k, v in config.items() if k in create_kwargs} - # create_args_keys = set(create_args.keys()) - # if not required_create_args.issubset(create_args_keys): - # raise ValueError(f"Required create args are missing: {required_create_args - create_args_keys}") - # if disallowed_create_args.intersection(create_args_keys): - # raise ValueError(f"Disallowed create args are present: {disallowed_create_args.intersection(create_args_keys)}") - # return create_args # TODO check types @@ -283,8 +315,17 @@ def convert_tools( if parameters is not None: ollama_properties = {} for prop_name, prop_schema in parameters["properties"].items(): + # Determine property type, checking "type" first, then "anyOf", defaulting to "string" + prop_type = prop_schema.get("type") + if prop_type is None and "anyOf" in prop_schema: + prop_type = next( + (opt.get("type") for opt in prop_schema["anyOf"] if opt.get("type") != "null"), + None, # Default to None if no non-null type found in anyOf + ) + prop_type = prop_type or "string" + ollama_properties[prop_name] = OllamaTool.Function.Parameters.Property( - type=prop_schema["type"], + type=prop_type, description=prop_schema["description"] if "description" in prop_schema else None, ) result.append( @@ -346,6 +387,66 @@ def normalize_stop_reason(stop_reason: str | None) -> FinishReasons: return KNOWN_STOP_MAPPINGS.get(stop_reason, "unknown") +# TODO: probably needs work +def count_tokens_ollama(messages: Sequence[LLMMessage], model: str, *, tools: Sequence[Tool | ToolSchema] = []) -> int: + try: + encoding = tiktoken.encoding_for_model(model) + except KeyError: + trace_logger.warning(f"Model {model} not found. Using cl100k_base encoding.") + encoding = tiktoken.get_encoding("cl100k_base") + tokens_per_message = 3 + num_tokens = 0 + + # Message tokens. + for message in messages: + num_tokens += tokens_per_message + ollama_message = to_ollama_type(message) + for ollama_message_part in ollama_message: + if isinstance(message.content, Image): + num_tokens += calculate_vision_tokens(message.content) + elif ollama_message_part.content is not None: + num_tokens += len(encoding.encode(ollama_message_part.content)) + # TODO: every model family has its own message sequence. + num_tokens += 3 # every reply is primed with <|start|>assistant<|message|> + + # Tool tokens. + ollama_tools = convert_tools(tools) + for tool in ollama_tools: + function = tool["function"] + tool_tokens = len(encoding.encode(function["name"])) + if "description" in function: + tool_tokens += len(encoding.encode(function["description"])) + tool_tokens -= 2 + if "parameters" in function: + parameters = function["parameters"] + if "properties" in parameters: + assert isinstance(parameters["properties"], dict) + for propertiesKey in parameters["properties"]: # pyright: ignore + assert isinstance(propertiesKey, str) + tool_tokens += len(encoding.encode(propertiesKey)) + v = parameters["properties"][propertiesKey] # pyright: ignore + for field in v: # pyright: ignore + if field == "type": + tool_tokens += 2 + tool_tokens += len(encoding.encode(v["type"])) # pyright: ignore + elif field == "description": + tool_tokens += 2 + tool_tokens += len(encoding.encode(v["description"])) # pyright: ignore + elif field == "enum": + tool_tokens -= 3 + for o in v["enum"]: # pyright: ignore + tool_tokens += 3 + tool_tokens += len(encoding.encode(o)) # pyright: ignore + else: + trace_logger.warning(f"Not supported field {field}") + tool_tokens += 11 + if len(parameters["properties"]) == 0: # pyright: ignore + tool_tokens -= 2 + num_tokens += tool_tokens + num_tokens += 12 + return num_tokens + + @dataclass class CreateParams: messages: Sequence[Message] @@ -413,6 +514,7 @@ def _process_create_args( self, messages: Sequence[LLMMessage], tools: Sequence[Tool | ToolSchema], + tool_choice: Tool | Literal["auto", "required", "none"], json_output: Optional[bool | type[BaseModel]], extra_create_args: Mapping[str, Any], ) -> CreateParams: @@ -483,7 +585,22 @@ def _process_create_args( if self.model_info["function_calling"] is False and len(tools) > 0: raise ValueError("Model does not support function calling and tools were provided") - converted_tools = convert_tools(tools) + converted_tools: List[OllamaTool] = [] + + # Handle tool_choice parameter in a way that is compatible with Ollama API. + if isinstance(tool_choice, Tool): + # If tool_choice is a Tool, convert it to OllamaTool. + converted_tools = convert_tools([tool_choice]) + elif tool_choice == "none": + # No tool choice, do not pass tools to the API. + converted_tools = [] + elif tool_choice == "required": + # Required tool choice, pass tools to the API. + converted_tools = convert_tools(tools) + if len(converted_tools) == 0: + raise ValueError("tool_choice 'required' specified but no tools provided") + else: + converted_tools = convert_tools(tools) return CreateParams( messages=ollama_messages, @@ -497,6 +614,7 @@ async def create( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -509,6 +627,7 @@ async def create( create_params = self._process_create_args( messages, tools, + tool_choice, json_output, extra_create_args, ) @@ -552,22 +671,10 @@ async def create( # Detect whether it is a function call or not. # We don't rely on choice.finish_reason as it is not always accurate, depending on the API used. content: Union[str, List[FunctionCall]] + thought: Optional[str] = None if result.message.tool_calls is not None: - # TODO: What are possible values for done_reason? - if result.done_reason != "tool_calls": - warnings.warn( - f"Finish reason mismatch: {result.done_reason} != tool_calls " - "when tool_calls are present. Finish reason may not be accurate. " - "This may be due to the API used that is not returning the correct finish reason.", - stacklevel=2, - ) - # TODO: Is this still an error condition? if result.message.content is not None and result.message.content != "": - warnings.warn( - "Both tool_calls and content are present in the message. " - "This is unexpected. content will be ignored, tool_calls will be used.", - stacklevel=2, - ) + thought = result.message.content # NOTE: If OAI response type changes, this will need to be updated content = [ FunctionCall( @@ -602,6 +709,7 @@ async def create( usage=usage, cached=False, logprobs=None, + thought=thought, ) self._total_usage = _add_usage(self._total_usage, usage) @@ -614,6 +722,7 @@ async def create_stream( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -626,6 +735,7 @@ async def create_stream( create_params = self._process_create_args( messages, tools, + tool_choice, json_output, extra_create_args, ) @@ -671,9 +781,8 @@ async def create_stream( content_chunks.append(chunk.message.content) if len(chunk.message.content) > 0: yield chunk.message.content - continue - # Otherwise, get tool calls + # Get tool calls if chunk.message.tool_calls is not None: full_tool_calls.extend( [ @@ -707,11 +816,17 @@ async def create_stream( else: prompt_tokens = 0 - if stop_reason == "function_call": - raise ValueError("Function calls are not supported in this context") - content: Union[str, List[FunctionCall]] - if len(content_chunks) > 1: + thought: Optional[str] = None + + if len(content_chunks) > 0 and len(full_tool_calls) > 0: + content = full_tool_calls + thought = "".join(content_chunks) + if chunk and chunk.eval_count: + completion_tokens = chunk.eval_count + else: + completion_tokens = 0 + elif len(content_chunks) > 1: content = "".join(content_chunks) if chunk and chunk.eval_count: completion_tokens = chunk.eval_count @@ -719,11 +834,6 @@ async def create_stream( completion_tokens = 0 else: completion_tokens = 0 - # TODO: fix assumption that dict values were added in order and actually order by int index - # for tool_call in full_tool_calls.values(): - # # value = json.dumps(tool_call) - # # completion_tokens += count_token(value, model=model) - # completion_tokens += 0 content = full_tool_calls usage = RequestUsage( @@ -737,6 +847,7 @@ async def create_stream( usage=usage, cached=False, logprobs=None, + thought=thought, ) # Emit the end event. @@ -762,65 +873,8 @@ def actual_usage(self) -> RequestUsage: def total_usage(self) -> RequestUsage: return self._total_usage - # TODO: probably needs work def count_tokens(self, messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = []) -> int: - model = self._create_args["model"] - try: - encoding = tiktoken.encoding_for_model(model) - except KeyError: - trace_logger.warning(f"Model {model} not found. Using cl100k_base encoding.") - encoding = tiktoken.get_encoding("cl100k_base") - tokens_per_message = 3 - num_tokens = 0 - - # Message tokens. - for message in messages: - num_tokens += tokens_per_message - ollama_message = to_ollama_type(message) - for ollama_message_part in ollama_message: - if isinstance(message.content, Image): - num_tokens += calculate_vision_tokens(message.content) - elif ollama_message_part.content is not None: - num_tokens += len(encoding.encode(ollama_message_part.content)) - # TODO: every model family has its own message sequence. - num_tokens += 3 # every reply is primed with <|start|>assistant<|message|> - - # Tool tokens. - ollama_tools = convert_tools(tools) - for tool in ollama_tools: - function = tool["function"] - tool_tokens = len(encoding.encode(function["name"])) - if "description" in function: - tool_tokens += len(encoding.encode(function["description"])) - tool_tokens -= 2 - if "parameters" in function: - parameters = function["parameters"] - if "properties" in parameters: - assert isinstance(parameters["properties"], dict) - for propertiesKey in parameters["properties"]: # pyright: ignore - assert isinstance(propertiesKey, str) - tool_tokens += len(encoding.encode(propertiesKey)) - v = parameters["properties"][propertiesKey] # pyright: ignore - for field in v: # pyright: ignore - if field == "type": - tool_tokens += 2 - tool_tokens += len(encoding.encode(v["type"])) # pyright: ignore - elif field == "description": - tool_tokens += 2 - tool_tokens += len(encoding.encode(v["description"])) # pyright: ignore - elif field == "enum": - tool_tokens -= 3 - for o in v["enum"]: # pyright: ignore - tool_tokens += 3 - tool_tokens += len(encoding.encode(o)) # pyright: ignore - else: - trace_logger.warning(f"Not supported field {field}") - tool_tokens += 11 - if len(parameters["properties"]) == 0: # pyright: ignore - tool_tokens -= 2 - num_tokens += tool_tokens - num_tokens += 12 - return num_tokens + return count_tokens_ollama(messages, self._create_args["model"], tools=tools) def remaining_tokens(self, messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = []) -> int: token_limit = _model_info.get_token_limit(self._create_args["model"]) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/__init__.py index 366ad831175e..2241f663af26 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/__init__.py @@ -1,4 +1,10 @@ -from ._openai_client import AzureOpenAIChatCompletionClient, BaseOpenAIChatCompletionClient, OpenAIChatCompletionClient +from . import _message_transform +from ._openai_client import ( + AZURE_OPENAI_USER_AGENT, + AzureOpenAIChatCompletionClient, + BaseOpenAIChatCompletionClient, + OpenAIChatCompletionClient, +) from .config import ( AzureOpenAIClientConfigurationConfigModel, BaseOpenAIClientConfigurationConfigModel, @@ -14,4 +20,6 @@ "OpenAIClientConfigurationConfigModel", "BaseOpenAIClientConfigurationConfigModel", "CreateArgumentsConfigModel", + "AZURE_OPENAI_USER_AGENT", + "_message_transform", ] diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_message_transform.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_message_transform.py new file mode 100644 index 000000000000..d21f9f95dfbf --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_message_transform.py @@ -0,0 +1,560 @@ +""" +# `_message_transform.py` Module-Level Documentation + +This document is a markdown-formatted version of the module-level docstring inserted into `_message_transform.py` as part of [PR #6063](https://github.com/microsoft/autogen/pull/6063). + +--- + +## AutoGen Modular Transformer Pipeline + +This module implements a modular and extensible message transformation pipeline +for converting `LLMMessage` instances into SDK-specific message formats +(e.g., OpenAI-style `ChatCompletionMessageParam`). + +--- + +### 📌 Background + +In previous versions of AutoGen, message adaptation was handled in ad-hoc ways, +scattered across model clients. This led to compatibility bugs and code duplication, +especially when supporting diverse models such as Gemini, Claude, or Anthropic SDKs. + +To address this, PR #6063 introduced a unified, composable transformer pipeline +that decouples message transformation logic from model SDK constructors. + +--- + +### 🎯 Key Concepts + +- **Transformer Function**: + Transforms a field (e.g., `content`, `name`, `role`) of an `LLMMessage` into a keyword argument. + +- **Transformer Pipeline**: + A sequence of transformer functions composed using `build_transformer_func`. + +- **Transformer Map**: + A dictionary mapping `LLMMessage` types (System, User, Assistant) to transformers for a specific model. + +- **Conditional Transformer**: + Chooses a pipeline dynamically based on message content or runtime conditions. + +--- + +### 🧪 Example: Basic Flow + +```python +from autogen_ext.models.openai._message_transform import get_transformer +from autogen.types import AssistantMessage + +llm_message = AssistantMessage(name="a", thought="Let's go!") +transformer = get_transformer("openai", "gpt-4", type(llm_message)) +sdk_message = transformer(llm_message, context={}) +print(sdk_message) +``` + +--- + +### 🧰 Example: Define Transformer Functions + +```python +def _set_role(role: str): + def fn(message, context): + return {"role": role} + + return fn + + +def _set_content_from_thought(message, context): + return {"content": message.thought or " "} + + +base_user_transformer_funcs = [_set_role("user"), _set_content_from_thought] +``` + +--- + +### 🛠️ Example: Build and Register Transformer Map + +```python +from autogen_ext.models.utils import build_transformer_func, register_transformer +from openai.types.chat import ChatCompletionUserMessageParam +from autogen.types import UserMessage, SystemMessage, AssistantMessage + +user_transformer = build_transformer_func( + funcs=base_user_transformer_funcs, message_param_func=ChatCompletionUserMessageParam +) + +MY_TRANSFORMER_MAP = {UserMessage: user_transformer, SystemMessage: ..., AssistantMessage: ...} + +register_transformer("openai", "mistral-7b", MY_TRANSFORMER_MAP) +``` + +--- + +### 🔁 Conditional Transformer Example + +```python +from autogen_ext.models.utils import build_conditional_transformer_func + + +def condition_func(message, context): + return "multimodal" if isinstance(message.content, dict) else "text" + + +user_transformers = { + "text": [_set_content_from_thought], + "multimodal": [_set_content_from_thought], # could be different logic +} + +message_param_funcs = { + "text": ChatCompletionUserMessageParam, + "multimodal": ChatCompletionUserMessageParam, +} + +conditional_user_transformer = build_conditional_transformer_func( + funcs_map=user_transformers, + message_param_func_map=message_param_funcs, + condition_func=condition_func, +) +``` + +--- + +### 📦 Design Principles + +- ✅ DRY and Composable +- ✅ Model-specific overrides without forking entire clients +- ✅ Explicit separation between transformation logic and SDK builders +- ✅ Future extensibility (e.g., Claude, Gemini, Alibaba) + +--- + +### 📎 Reference + +- Introduced in: [PR #6063](https://github.com/microsoft/autogen/pull/6063) +""" + +from typing import Any, Callable, Dict, List, cast, get_args + +from autogen_core import ( + FunctionCall, + Image, +) +from autogen_core.models import ( + AssistantMessage, + FunctionExecutionResultMessage, + LLMMessage, + ModelFamily, + SystemMessage, + UserMessage, +) +from openai.types.chat import ( + ChatCompletionAssistantMessageParam, + ChatCompletionContentPartImageParam, + ChatCompletionContentPartParam, + ChatCompletionContentPartTextParam, + ChatCompletionMessageToolCallParam, + ChatCompletionSystemMessageParam, + ChatCompletionToolMessageParam, + ChatCompletionUserMessageParam, +) + +from ._transformation import ( + LLMMessageContent, + TransformerMap, + TrasformerReturnType, + build_conditional_transformer_func, + build_transformer_func, + register_transformer, +) +from ._utils import assert_valid_name + +EMPTY: Dict[str, Any] = {} + + +def func_call_to_oai(message: FunctionCall) -> ChatCompletionMessageToolCallParam: + return ChatCompletionMessageToolCallParam( + id=message.id, + function={ + "arguments": message.arguments, + "name": message.name, + }, + type="function", + ) + + +# ===Mini Transformers=== +def _assert_valid_name(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, None]: + assert isinstance(message, (UserMessage, AssistantMessage)) + assert_valid_name(message.source) + return EMPTY + + +def _set_role(role: str) -> Callable[[LLMMessage, Dict[str, Any]], Dict[str, str]]: + def inner(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, str]: + return {"role": role} + + return inner + + +def _set_name(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, Any]: + assert isinstance(message, (UserMessage, AssistantMessage)) + assert_valid_name(message.source) + # Check if name should be included in message + if context.get("include_name_in_message", True): + return {"name": message.source} + else: + return EMPTY + + +def _set_content_direct(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, LLMMessageContent]: + return {"content": message.content} + + +def _set_prepend_text_content(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, str]: + assert isinstance(message, (UserMessage, AssistantMessage)) + assert isinstance(message.content, str) + prepend = context.get("prepend_name", False) + prefix = f"{message.source} said:\n" if prepend else "" + return {"content": prefix + message.content} + + +def _set_multimodal_content( + message: LLMMessage, context: Dict[str, Any] +) -> Dict[str, List[ChatCompletionContentPartParam]]: + assert isinstance(message, (UserMessage, AssistantMessage)) + prepend = context.get("prepend_name", False) + parts: List[ChatCompletionContentPartParam] = [] + + for idx, part in enumerate(message.content): + if isinstance(part, str): + # If prepend, Append the name to the first text part + text = f"{message.source} said:\n" + part if prepend and idx == 0 else part + parts.append(ChatCompletionContentPartTextParam(type="text", text=text)) + elif isinstance(part, Image): + # TODO: support url based images + # TODO: support specifying details + parts.append(cast(ChatCompletionContentPartImageParam, part.to_openai_format())) + else: + raise ValueError(f"Unknown content part: {part}") + + return {"content": parts} + + +def _set_tool_calls( + message: LLMMessage, context: Dict[str, Any] +) -> Dict[str, List[ChatCompletionMessageToolCallParam]]: + assert isinstance(message.content, list) + assert isinstance(message, AssistantMessage) + return { + "tool_calls": [func_call_to_oai(x) for x in message.content], + } + + +def _set_thought_as_content(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, str | None]: + assert isinstance(message, AssistantMessage) + return {"content": message.thought} + + +def _set_thought_as_content_gemini(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, str | None]: + assert isinstance(message, AssistantMessage) + return {"content": message.thought or " "} + + +def _set_empty_to_whitespace(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, LLMMessageContent]: + return {"content": message.content or " "} + + +def _set_pass_message_when_whitespace(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, bool]: + if isinstance(message.content, str) and (message.content.isspace() or not message.content): + return {"pass_message": True} + return {} + + +def _set_null_content_for_tool_calls(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, None]: + """Set content to null for tool calls without thought. Required by OpenAI API.""" + assert isinstance(message, AssistantMessage) + return {"content": None} + + +# === Base Transformers list === +base_system_message_transformers: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = [ + _set_content_direct, + _set_role("system"), +] + +base_user_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = [ + _assert_valid_name, + _set_role("user"), +] + +base_assistant_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = [ + _assert_valid_name, + _set_role("assistant"), +] + + +# === Transformers list === +system_message_transformers: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = ( + base_system_message_transformers +) + +single_user_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = ( + base_user_transformer_funcs + + [ + _set_name, + _set_prepend_text_content, + ] +) + +multimodal_user_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = ( + base_user_transformer_funcs + + [ + _set_name, + _set_multimodal_content, + ] +) + +single_assistant_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = ( + base_assistant_transformer_funcs + + [ + _set_content_direct, + ] +) + +tools_assistant_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = ( + base_assistant_transformer_funcs + + [ + _set_tool_calls, + _set_null_content_for_tool_calls, + ] +) + +thought_assistant_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = ( + base_assistant_transformer_funcs + + [ + _set_tool_calls, + _set_thought_as_content, + ] +) + +thought_assistant_transformer_funcs_gemini: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = ( + base_assistant_transformer_funcs + + [ + _set_tool_calls, + _set_thought_as_content_gemini, + ] +) + + +# === Specific message param functions === +single_user_transformer_funcs_mistral: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = ( + base_user_transformer_funcs + + [ + _set_prepend_text_content, + ] +) + +multimodal_user_transformer_funcs_mistral: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = ( + base_user_transformer_funcs + + [ + _set_multimodal_content, + ] +) + + +# === Transformer maps === +user_transformer_funcs: Dict[str, List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]]] = { + "text": single_user_transformer_funcs, + "multimodal": multimodal_user_transformer_funcs, +} +user_transformer_constructors: Dict[str, Callable[..., Any]] = { + "text": ChatCompletionUserMessageParam, + "multimodal": ChatCompletionUserMessageParam, +} + + +def user_condition(message: LLMMessage, context: Dict[str, Any]) -> str: + if isinstance(message.content, str): + return "text" + else: + return "multimodal" + + +assistant_transformer_funcs: Dict[str, List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]]] = { + "text": single_assistant_transformer_funcs, + "tools": tools_assistant_transformer_funcs, + "thought": thought_assistant_transformer_funcs, +} + + +assistant_transformer_constructors: Dict[str, Callable[..., Any]] = { + "text": ChatCompletionAssistantMessageParam, + "tools": ChatCompletionAssistantMessageParam, + "thought": ChatCompletionAssistantMessageParam, +} + + +def assistant_condition(message: LLMMessage, context: Dict[str, Any]) -> str: + assert isinstance(message, AssistantMessage) + if isinstance(message.content, list): + if message.thought is not None: + return "thought" + else: + return "tools" + else: + return "text" + + +user_transformer_funcs_gemini: Dict[str, List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]]] = { + "text": single_user_transformer_funcs + [_set_empty_to_whitespace], + "multimodal": multimodal_user_transformer_funcs, +} + + +assistant_transformer_funcs_gemini: Dict[str, List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]]] = { + "text": single_assistant_transformer_funcs + [_set_empty_to_whitespace], + "tools": tools_assistant_transformer_funcs, # that case, message.content is a list of FunctionCall + "thought": thought_assistant_transformer_funcs_gemini, # that case, message.content is a list of FunctionCall +} + + +user_transformer_funcs_claude: Dict[str, List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]]] = { + "text": single_user_transformer_funcs + [_set_pass_message_when_whitespace], + "multimodal": multimodal_user_transformer_funcs + [_set_pass_message_when_whitespace], +} + + +assistant_transformer_funcs_claude: Dict[str, List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]]] = { + "text": single_assistant_transformer_funcs + [_set_pass_message_when_whitespace], + "tools": tools_assistant_transformer_funcs, # that case, message.content is a list of FunctionCall + "thought": thought_assistant_transformer_funcs_gemini, # that case, message.content is a list of FunctionCall +} + + +user_transformer_funcs_mistral: Dict[str, List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]]] = { + "text": single_user_transformer_funcs_mistral, + "multimodal": multimodal_user_transformer_funcs_mistral, +} + + +def function_execution_result_message(message: LLMMessage, context: Dict[str, Any]) -> TrasformerReturnType: + assert isinstance(message, FunctionExecutionResultMessage) + return [ + ChatCompletionToolMessageParam(content=x.content, role="tool", tool_call_id=x.call_id) for x in message.content + ] + + +# === Transformers === + +__BASE_TRANSFORMER_MAP: TransformerMap = { + SystemMessage: build_transformer_func( + funcs=system_message_transformers, + message_param_func=ChatCompletionSystemMessageParam, + ), + UserMessage: build_conditional_transformer_func( + funcs_map=user_transformer_funcs, + message_param_func_map=user_transformer_constructors, + condition_func=user_condition, + ), + AssistantMessage: build_conditional_transformer_func( + funcs_map=assistant_transformer_funcs, + message_param_func_map=assistant_transformer_constructors, + condition_func=assistant_condition, + ), + FunctionExecutionResultMessage: function_execution_result_message, +} + +__GEMINI_TRANSFORMER_MAP: TransformerMap = { + SystemMessage: build_transformer_func( + funcs=system_message_transformers + [_set_empty_to_whitespace], + message_param_func=ChatCompletionSystemMessageParam, + ), + UserMessage: build_conditional_transformer_func( + funcs_map=user_transformer_funcs_gemini, + message_param_func_map=user_transformer_constructors, + condition_func=user_condition, + ), + AssistantMessage: build_conditional_transformer_func( + funcs_map=assistant_transformer_funcs_gemini, + message_param_func_map=assistant_transformer_constructors, + condition_func=assistant_condition, + ), + FunctionExecutionResultMessage: function_execution_result_message, +} + +__CLAUDE_TRANSFORMER_MAP: TransformerMap = { + SystemMessage: build_transformer_func( + funcs=system_message_transformers + [_set_empty_to_whitespace], + message_param_func=ChatCompletionSystemMessageParam, + ), + UserMessage: build_conditional_transformer_func( + funcs_map=user_transformer_funcs_claude, + message_param_func_map=user_transformer_constructors, + condition_func=user_condition, + ), + AssistantMessage: build_conditional_transformer_func( + funcs_map=assistant_transformer_funcs_claude, + message_param_func_map=assistant_transformer_constructors, + condition_func=assistant_condition, + ), + FunctionExecutionResultMessage: function_execution_result_message, +} + +__MISTRAL_TRANSFORMER_MAP: TransformerMap = { + SystemMessage: build_transformer_func( + funcs=system_message_transformers + [_set_empty_to_whitespace], + message_param_func=ChatCompletionSystemMessageParam, + ), + UserMessage: build_conditional_transformer_func( + funcs_map=user_transformer_funcs_mistral, + message_param_func_map=user_transformer_constructors, + condition_func=user_condition, + ), + AssistantMessage: build_conditional_transformer_func( + funcs_map=assistant_transformer_funcs, + message_param_func_map=assistant_transformer_constructors, + condition_func=assistant_condition, + ), + FunctionExecutionResultMessage: function_execution_result_message, +} + + +# set openai models to use the transformer map +total_models = get_args(ModelFamily.ANY) +__openai_models = [model for model in total_models if ModelFamily.is_openai(model)] + +__claude_models = [model for model in total_models if ModelFamily.is_claude(model)] + +__gemini_models = [model for model in total_models if ModelFamily.is_gemini(model)] + +__llama_models = [model for model in total_models if ModelFamily.is_llama(model)] + +__unknown_models = list( + set(total_models) - set(__openai_models) - set(__claude_models) - set(__gemini_models) - set(__llama_models) +) +__mistral_models = [model for model in total_models if ModelFamily.is_mistral(model)] + +__unknown_models = list( + set(total_models) - set(__openai_models) - set(__claude_models) - set(__gemini_models) - set(__mistral_models) +) + +for model in __openai_models: + register_transformer("openai", model, __BASE_TRANSFORMER_MAP) + +for model in __claude_models: + register_transformer("openai", model, __CLAUDE_TRANSFORMER_MAP) + +for model in __gemini_models: + register_transformer("openai", model, __GEMINI_TRANSFORMER_MAP) + +for model in __llama_models: + register_transformer("openai", model, __BASE_TRANSFORMER_MAP) + +for model in __mistral_models: + register_transformer("openai", model, __MISTRAL_TRANSFORMER_MAP) + +for model in __unknown_models: + register_transformer("openai", model, __BASE_TRANSFORMER_MAP) + +register_transformer("openai", "default", __BASE_TRANSFORMER_MAP) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_model_info.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_model_info.py index 20a4177ff3c9..6306fba941cf 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_model_info.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_model_info.py @@ -1,14 +1,26 @@ +import logging from typing import Dict +from autogen_core import EVENT_LOGGER_NAME, TRACE_LOGGER_NAME from autogen_core.models import ModelFamily, ModelInfo +logger = logging.getLogger(EVENT_LOGGER_NAME) +trace_logger = logging.getLogger(TRACE_LOGGER_NAME) + # Based on: https://platform.openai.com/docs/models/continuous-model-upgrades # This is a moving target, so correctness is checked by the model value returned by openai against expected values at runtime`` _MODEL_POINTERS = { + # OpenAI models + "o4-mini": "o4-mini-2025-04-16", + "o3": "o3-2025-04-16", "o3-mini": "o3-mini-2025-01-31", "o1": "o1-2024-12-17", "o1-preview": "o1-preview-2024-09-12", "o1-mini": "o1-mini-2024-09-12", + "gpt-4.1": "gpt-4.1-2025-04-14", + "gpt-4.1-mini": "gpt-4.1-mini-2025-04-14", + "gpt-4.1-nano": "gpt-4.1-nano-2025-04-14", + "gpt-4.5-preview": "gpt-4.5-preview-2025-02-27", "gpt-4o": "gpt-4o-2024-08-06", "gpt-4o-mini": "gpt-4o-mini-2024-07-18", "gpt-4-turbo": "gpt-4-turbo-2024-04-09", @@ -17,29 +29,78 @@ "gpt-4-32k": "gpt-4-32k-0613", "gpt-3.5-turbo": "gpt-3.5-turbo-0125", "gpt-3.5-turbo-16k": "gpt-3.5-turbo-16k-0613", + # Anthropic models + "claude-3-haiku": "claude-3-haiku-20240307", + "claude-3-sonnet": "claude-3-sonnet-20240229", + "claude-3-opus": "claude-3-opus-20240229", + "claude-3-5-haiku": "claude-3-5-haiku-20241022", + "claude-3-5-sonnet": "claude-3-5-sonnet-20241022", + "claude-3-7-sonnet": "claude-3-7-sonnet-20250219", + "claude-4-sonnet": "claude-sonnet-4-20250514", + "claude-4-opus": "claude-opus-4-20250514", + # Llama models + "llama-3.3-8b": "Llama-3.3-8B-Instruct", + "llama-3.3-70b": "Llama-3.3-70B-Instruct", + "llama-4-scout": "Llama-4-Scout-17B-16E-Instruct-FP8", + "llama-4-maverick": "Llama-4-Maverick-17B-128E-Instruct-FP8", } _MODEL_INFO: Dict[str, ModelInfo] = { + "gpt-4o-mini-search-preview-2025-03-11": { + "vision": False, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GPT_4O, + "structured_output": True, + "multiple_system_messages": True, + }, + "gpt-4o-search-preview-2025-03-11": { + "vision": False, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GPT_4O, + "structured_output": True, + "multiple_system_messages": True, + }, + "o4-mini-2025-04-16": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.O4, + "structured_output": True, + "multiple_system_messages": True, + }, + "o3-2025-04-16": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.O3, + "structured_output": True, + "multiple_system_messages": True, + }, "o3-mini-2025-01-31": { "vision": False, "function_calling": True, "json_output": True, "family": ModelFamily.O3, "structured_output": True, + "multiple_system_messages": True, }, "o1-2024-12-17": { "vision": False, - "function_calling": False, + "function_calling": True, "json_output": False, "family": ModelFamily.O1, "structured_output": True, + "multiple_system_messages": True, }, "o1-preview-2024-09-12": { "vision": False, - "function_calling": False, + "function_calling": True, "json_output": False, "family": ModelFamily.O1, "structured_output": True, + "multiple_system_messages": True, }, "o1-mini-2024-09-12": { "vision": False, @@ -47,6 +108,39 @@ "json_output": False, "family": ModelFamily.O1, "structured_output": False, + "multiple_system_messages": True, + }, + "gpt-4.1-2025-04-14": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GPT_41, + "structured_output": True, + "multiple_system_messages": True, + }, + "gpt-4.1-mini-2025-04-14": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GPT_41, + "structured_output": True, + "multiple_system_messages": True, + }, + "gpt-4.1-nano-2025-04-14": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GPT_41, + "structured_output": True, + "multiple_system_messages": True, + }, + "gpt-4.5-preview-2025-02-27": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GPT_45, + "structured_output": True, + "multiple_system_messages": True, }, "gpt-4o-2024-11-20": { "vision": True, @@ -54,6 +148,7 @@ "json_output": True, "family": ModelFamily.GPT_4O, "structured_output": True, + "multiple_system_messages": True, }, "gpt-4o-2024-08-06": { "vision": True, @@ -61,6 +156,7 @@ "json_output": True, "family": ModelFamily.GPT_4O, "structured_output": True, + "multiple_system_messages": True, }, "gpt-4o-2024-05-13": { "vision": True, @@ -68,6 +164,7 @@ "json_output": True, "family": ModelFamily.GPT_4O, "structured_output": False, + "multiple_system_messages": True, }, "gpt-4o-mini-2024-07-18": { "vision": True, @@ -75,6 +172,7 @@ "json_output": True, "family": ModelFamily.GPT_4O, "structured_output": True, + "multiple_system_messages": True, }, "gpt-4-turbo-2024-04-09": { "vision": True, @@ -82,6 +180,7 @@ "json_output": True, "family": ModelFamily.GPT_4, "structured_output": False, + "multiple_system_messages": True, }, "gpt-4-0125-preview": { "vision": False, @@ -89,6 +188,7 @@ "json_output": True, "family": ModelFamily.GPT_4, "structured_output": False, + "multiple_system_messages": True, }, "gpt-4-1106-preview": { "vision": False, @@ -96,6 +196,7 @@ "json_output": True, "family": ModelFamily.GPT_4, "structured_output": False, + "multiple_system_messages": True, }, "gpt-4-1106-vision-preview": { "vision": True, @@ -103,6 +204,7 @@ "json_output": False, "family": ModelFamily.GPT_4, "structured_output": False, + "multiple_system_messages": True, }, "gpt-4-0613": { "vision": False, @@ -110,6 +212,7 @@ "json_output": True, "family": ModelFamily.GPT_4, "structured_output": False, + "multiple_system_messages": True, }, "gpt-4-32k-0613": { "vision": False, @@ -117,6 +220,7 @@ "json_output": True, "family": ModelFamily.GPT_4, "structured_output": False, + "multiple_system_messages": True, }, "gpt-3.5-turbo-0125": { "vision": False, @@ -124,6 +228,7 @@ "json_output": True, "family": ModelFamily.GPT_35, "structured_output": False, + "multiple_system_messages": True, }, "gpt-3.5-turbo-1106": { "vision": False, @@ -131,6 +236,7 @@ "json_output": True, "family": ModelFamily.GPT_35, "structured_output": False, + "multiple_system_messages": True, }, "gpt-3.5-turbo-instruct": { "vision": False, @@ -138,6 +244,7 @@ "json_output": True, "family": ModelFamily.GPT_35, "structured_output": False, + "multiple_system_messages": True, }, "gpt-3.5-turbo-0613": { "vision": False, @@ -145,6 +252,7 @@ "json_output": True, "family": ModelFamily.GPT_35, "structured_output": False, + "multiple_system_messages": True, }, "gpt-3.5-turbo-16k-0613": { "vision": False, @@ -152,6 +260,7 @@ "json_output": True, "family": ModelFamily.GPT_35, "structured_output": False, + "multiple_system_messages": True, }, "gemini-1.5-flash": { "vision": True, @@ -159,6 +268,7 @@ "json_output": True, "family": ModelFamily.GEMINI_1_5_FLASH, "structured_output": True, + "multiple_system_messages": False, }, "gemini-1.5-flash-8b": { "vision": True, @@ -166,6 +276,7 @@ "json_output": True, "family": ModelFamily.GEMINI_1_5_FLASH, "structured_output": True, + "multiple_system_messages": False, }, "gemini-1.5-pro": { "vision": True, @@ -173,6 +284,7 @@ "json_output": True, "family": ModelFamily.GEMINI_1_5_PRO, "structured_output": True, + "multiple_system_messages": False, }, "gemini-2.0-flash": { "vision": True, @@ -180,6 +292,7 @@ "json_output": True, "family": ModelFamily.GEMINI_2_0_FLASH, "structured_output": True, + "multiple_system_messages": False, }, "gemini-2.0-flash-lite-preview-02-05": { "vision": True, @@ -187,14 +300,133 @@ "json_output": True, "family": ModelFamily.GEMINI_2_0_FLASH, "structured_output": True, + "multiple_system_messages": False, + }, + "gemini-2.5-pro-preview-03-25": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GEMINI_2_5_PRO, + "structured_output": True, + "multiple_system_messages": False, + }, + "gemini-2.5-flash": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GEMINI_2_5_FLASH, + "structured_output": True, + "multiple_system_messages": False, + }, + "claude-3-haiku-20240307": { + "vision": True, + "function_calling": True, + "json_output": False, # Update this when Anthropic supports structured output + "family": ModelFamily.CLAUDE_3_HAIKU, + "structured_output": False, + "multiple_system_messages": True, + }, + "claude-3-sonnet-20240229": { + "vision": True, + "function_calling": True, + "json_output": False, # Update this when Anthropic supports structured output + "family": ModelFamily.CLAUDE_3_SONNET, + "structured_output": False, + "multiple_system_messages": True, + }, + "claude-3-opus-20240229": { + "vision": True, + "function_calling": True, + "json_output": False, # Update this when Anthropic supports structured output + "family": ModelFamily.CLAUDE_3_OPUS, + "structured_output": False, + "multiple_system_messages": True, + }, + "claude-3-5-haiku-20241022": { + "vision": True, + "function_calling": True, + "json_output": False, # Update this when Anthropic supports structured output + "family": ModelFamily.CLAUDE_3_5_HAIKU, + "structured_output": False, + "multiple_system_messages": True, + }, + "claude-3-5-sonnet-20241022": { + "vision": True, + "function_calling": True, + "json_output": False, # Update this when Anthropic supports structured output + "family": ModelFamily.CLAUDE_3_5_SONNET, + "structured_output": False, + "multiple_system_messages": True, + }, + "claude-3-7-sonnet-20250219": { + "vision": True, + "function_calling": True, + "json_output": False, # Update this when Anthropic supports structured output + "family": ModelFamily.CLAUDE_3_7_SONNET, + "structured_output": False, + "multiple_system_messages": True, + }, + "claude-sonnet-4-20250514": { + "vision": True, + "function_calling": True, + "json_output": False, # Update this when Anthropic supports structured output + "family": ModelFamily.CLAUDE_4_SONNET, + "structured_output": False, + "multiple_system_messages": True, + }, + "claude-opus-4-20250514": { + "vision": True, + "function_calling": True, + "json_output": False, # Update this when Anthropic supports structured output + "family": ModelFamily.CLAUDE_4_OPUS, + "structured_output": False, + "multiple_system_messages": True, + }, + "Llama-3.3-8B-Instruct": { + "vision": False, + "function_calling": True, + "json_output": True, + "family": ModelFamily.LLAMA_3_3_8B, + "structured_output": False, + "multiple_system_messages": True, + }, + "Llama-3.3-70B-Instruct": { + "vision": False, + "function_calling": True, + "json_output": True, + "family": ModelFamily.LLAMA_3_3_70B, + "structured_output": False, + "multiple_system_messages": True, + }, + "Llama-4-Scout-17B-16E-Instruct-FP8": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.LLAMA_4_SCOUT, + "structured_output": True, + "multiple_system_messages": True, + }, + "Llama-4-Maverick-17B-128E-Instruct-FP8": { + "vision": True, + "function_calling": True, + "json_output": True, + "family": ModelFamily.LLAMA_4_MAVERICK, + "structured_output": True, + "multiple_system_messages": True, }, } _MODEL_TOKEN_LIMITS: Dict[str, int] = { + "o4-mini-2025-04-16": 200000, + "o3-2025-04-16": 200000, "o3-mini-2025-01-31": 200000, "o1-2024-12-17": 200000, "o1-preview-2024-09-12": 128000, "o1-mini-2024-09-12": 128000, + "gpt-4.1-2025-04-14": 1047576, + "gpt-4.1-mini-2025-04-14": 1047576, + "gpt-4.1-nano-2025-04-14": 1047576, + "gpt-4.5-preview-2025-02-27": 128000, "gpt-4o-2024-11-20": 128000, "gpt-4o-2024-08-06": 128000, "gpt-4o-2024-05-13": 128000, @@ -214,9 +446,26 @@ "gemini-1.5-flash-8b": 1048576, "gemini-1.5-pro": 2097152, "gemini-2.0-flash": 1048576, + "gemini-2.0-flash-lite-preview-02-05": 1048576, + "gemini-2.5-pro-preview-03-25": 2097152, + "gemini-2.5-flash": 1048576, + "claude-3-haiku-20240307": 50000, + "claude-3-sonnet-20240229": 200000, + "claude-3-opus-20240229": 200000, + "claude-3-5-haiku-20241022": 200000, + "claude-3-5-sonnet-20241022": 200000, + "claude-3-7-sonnet-20250219": 200000, + "claude-sonnet-4-20250514": 200000, + "claude-opus-4-20250514": 200000, + "Llama-3.3-8B-Instruct": 128000, + "Llama-3.3-70B-Instruct": 128000, + "Llama-4-Scout-17B-16E-Instruct-FP8": 128000, + "Llama-4-Maverick-17B-128E-Instruct-FP8": 128000, } GEMINI_OPENAI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/" +ANTHROPIC_OPENAI_BASE_URL = "https://api.anthropic.com/v1/" +LLAMA_API_BASE_URL = "https://api.llama.com/compat/v1/" def resolve_model(model: str) -> str: @@ -226,8 +475,24 @@ def resolve_model(model: str) -> str: def get_info(model: str) -> ModelInfo: + # If call it, that mean is that the config does not have cumstom model_info resolved_model = resolve_model(model) - return _MODEL_INFO[resolved_model] + model_info: ModelInfo = _MODEL_INFO.get( + resolved_model, + { + "vision": False, + "function_calling": False, + "json_output": False, + "family": "FAILED", + "structured_output": False, + }, + ) + if model_info.get("family") == "FAILED": + raise ValueError("model_info is required when model name is not a valid OpenAI model") + if model_info.get("family") == ModelFamily.UNKNOWN: + trace_logger.warning(f"Model info not found for model: {model}") + + return model_info def get_token_limit(model: str) -> int: diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py index efdb22cd1fa3..02b8d911a31a 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py @@ -8,11 +8,14 @@ import warnings from asyncio import Task from dataclasses import dataclass +from importlib.metadata import PackageNotFoundError, version from typing import ( Any, AsyncGenerator, + Callable, Dict, List, + Literal, Mapping, Optional, Sequence, @@ -37,7 +40,6 @@ ChatCompletionClient, ChatCompletionTokenLogprob, CreateResult, - FunctionExecutionResultMessage, LLMMessage, ModelCapabilities, # type: ignore ModelFamily, @@ -52,18 +54,11 @@ from openai import NOT_GIVEN, AsyncAzureOpenAI, AsyncOpenAI from openai.types.chat import ( ChatCompletion, - ChatCompletionAssistantMessageParam, ChatCompletionChunk, - ChatCompletionContentPartImageParam, ChatCompletionContentPartParam, - ChatCompletionContentPartTextParam, ChatCompletionMessageParam, - ChatCompletionMessageToolCallParam, ChatCompletionRole, - ChatCompletionSystemMessageParam, - ChatCompletionToolMessageParam, ChatCompletionToolParam, - ChatCompletionUserMessageParam, ParsedChatCompletion, ParsedChoice, completion_create_params, @@ -81,6 +76,10 @@ from .._utils.normalize_stop_reason import normalize_stop_reason from .._utils.parse_r1_content import parse_r1_content from . import _model_info +from ._transformation import ( + get_transformer, +) +from ._utils import assert_valid_name from .config import ( AzureOpenAIClientConfiguration, AzureOpenAIClientConfigurationConfigModel, @@ -101,12 +100,31 @@ disallowed_create_args = set(["stream", "messages", "function_call", "functions", "n"]) required_create_args: Set[str] = set(["model"]) +USER_AGENT_HEADER_NAME = "User-Agent" + +try: + version_info = version("autogen-ext") +except PackageNotFoundError: + version_info = "dev" +AZURE_OPENAI_USER_AGENT = f"autogen-python/{version_info}" + def _azure_openai_client_from_config(config: Mapping[str, Any]) -> AsyncAzureOpenAI: # Take a copy copied_config = dict(config).copy() # Shave down the config to just the AzureOpenAIChatCompletionClient kwargs azure_config = {k: v for k, v in copied_config.items() if k in aopenai_init_kwargs} + + DEFAULT_HEADERS_KEY = "default_headers" + if DEFAULT_HEADERS_KEY not in azure_config: + azure_config[DEFAULT_HEADERS_KEY] = {} + + azure_config[DEFAULT_HEADERS_KEY][USER_AGENT_HEADER_NAME] = ( + f"{AZURE_OPENAI_USER_AGENT} {azure_config[DEFAULT_HEADERS_KEY][USER_AGENT_HEADER_NAME]}" + if USER_AGENT_HEADER_NAME in azure_config[DEFAULT_HEADERS_KEY] + else AZURE_OPENAI_USER_AGENT + ) + return AsyncAzureOpenAI(**azure_config) @@ -144,105 +162,27 @@ def type_to_role(message: LLMMessage) -> ChatCompletionRole: return "tool" -def user_message_to_oai(message: UserMessage, prepend_name: bool = False) -> ChatCompletionUserMessageParam: - assert_valid_name(message.source) - if isinstance(message.content, str): - return ChatCompletionUserMessageParam( - content=(f"{message.source} said:\n" if prepend_name else "") + message.content, - role="user", - name=message.source, - ) - else: - parts: List[ChatCompletionContentPartParam] = [] - for part in message.content: - if isinstance(part, str): - if prepend_name: - # Append the name to the first text part - oai_part = ChatCompletionContentPartTextParam( - text=f"{message.source} said:\n" + part, - type="text", - ) - prepend_name = False - else: - oai_part = ChatCompletionContentPartTextParam( - text=part, - type="text", - ) - parts.append(oai_part) - elif isinstance(part, Image): - # TODO: support url based images - # TODO: support specifying details - parts.append(cast(ChatCompletionContentPartImageParam, part.to_openai_format())) - else: - raise ValueError(f"Unknown content type: {part}") - return ChatCompletionUserMessageParam( - content=parts, - role="user", - name=message.source, - ) - - -def system_message_to_oai(message: SystemMessage) -> ChatCompletionSystemMessageParam: - return ChatCompletionSystemMessageParam( - content=message.content, - role="system", +def to_oai_type( + message: LLMMessage, + prepend_name: bool = False, + model: str = "unknown", + model_family: str = ModelFamily.UNKNOWN, + include_name_in_message: bool = True, +) -> Sequence[ChatCompletionMessageParam]: + context = { + "prepend_name": prepend_name, + "include_name_in_message": include_name_in_message, + } + transformers = get_transformer("openai", model, model_family) + + def raise_value_error(message: LLMMessage, context: Dict[str, Any]) -> Sequence[ChatCompletionMessageParam]: + raise ValueError(f"Unknown message type: {type(message)}") + + transformer: Callable[[LLMMessage, Dict[str, Any]], Sequence[ChatCompletionMessageParam]] = transformers.get( + type(message), raise_value_error ) - - -def func_call_to_oai(message: FunctionCall) -> ChatCompletionMessageToolCallParam: - return ChatCompletionMessageToolCallParam( - id=message.id, - function={ - "arguments": message.arguments, - "name": message.name, - }, - type="function", - ) - - -def tool_message_to_oai( - message: FunctionExecutionResultMessage, -) -> Sequence[ChatCompletionToolMessageParam]: - return [ - ChatCompletionToolMessageParam(content=x.content, role="tool", tool_call_id=x.call_id) for x in message.content - ] - - -def assistant_message_to_oai( - message: AssistantMessage, -) -> ChatCompletionAssistantMessageParam: - assert_valid_name(message.source) - if isinstance(message.content, list): - if message.thought is not None: - return ChatCompletionAssistantMessageParam( - content=message.thought, - tool_calls=[func_call_to_oai(x) for x in message.content], - role="assistant", - name=message.source, - ) - else: - return ChatCompletionAssistantMessageParam( - tool_calls=[func_call_to_oai(x) for x in message.content], - role="assistant", - name=message.source, - ) - else: - return ChatCompletionAssistantMessageParam( - content=message.content, - role="assistant", - name=message.source, - ) - - -def to_oai_type(message: LLMMessage, prepend_name: bool = False) -> Sequence[ChatCompletionMessageParam]: - if isinstance(message, SystemMessage): - return [system_message_to_oai(message)] - elif isinstance(message, UserMessage): - return [user_message_to_oai(message, prepend_name)] - elif isinstance(message, AssistantMessage): - return [assistant_message_to_oai(message)] - else: - return tool_message_to_oai(message) + result = transformer(message, context) + return result def calculate_vision_tokens(image: Image, detail: str = "auto") -> int: @@ -331,6 +271,31 @@ def convert_tools( return result +def convert_tool_choice(tool_choice: Tool | Literal["auto", "required", "none"]) -> Any: + """Convert tool_choice parameter to OpenAI API format. + + Args: + tool_choice: A single Tool object to force the model to use, "auto" to let the model choose any available tool, "required" to force tool usage, or "none" to disable tool usage. + + Returns: + OpenAI API compatible tool_choice value or None if not specified. + """ + if tool_choice == "none": + return "none" + + if tool_choice == "auto": + return "auto" + + if tool_choice == "required": + return "required" + + # Must be a Tool object + if isinstance(tool_choice, Tool): + return {"type": "function", "function": {"name": tool_choice.schema["name"]}} + else: + raise ValueError(f"tool_choice must be a Tool object, 'auto', 'required', or 'none', got {type(tool_choice)}") + + def normalize_name(name: str) -> str: """ LLMs sometimes ask functions while ignoring their own format requirements, this function should be used to replace invalid characters with "_". @@ -340,17 +305,107 @@ def normalize_name(name: str) -> str: return re.sub(r"[^a-zA-Z0-9_-]", "_", name)[:64] -def assert_valid_name(name: str) -> str: - """ - Ensure that configured names are valid, raises ValueError if not. - - For munging LLM responses use _normalize_name to ensure LLM specified names don't break the API. - """ - if not re.match(r"^[a-zA-Z0-9_-]+$", name): - raise ValueError(f"Invalid name: {name}. Only letters, numbers, '_' and '-' are allowed.") - if len(name) > 64: - raise ValueError(f"Invalid name: {name}. Name must be less than 64 characters.") - return name +def count_tokens_openai( + messages: Sequence[LLMMessage], + model: str, + *, + add_name_prefixes: bool = False, + tools: Sequence[Tool | ToolSchema] = [], + model_family: str = ModelFamily.UNKNOWN, + include_name_in_message: bool = True, +) -> int: + try: + encoding = tiktoken.encoding_for_model(model) + except KeyError: + trace_logger.warning(f"Model {model} not found. Using cl100k_base encoding.") + encoding = tiktoken.get_encoding("cl100k_base") + tokens_per_message = 3 + tokens_per_name = 1 + num_tokens = 0 + + # Message tokens. + for message in messages: + num_tokens += tokens_per_message + oai_message = to_oai_type( + message, + prepend_name=add_name_prefixes, + model=model, + model_family=model_family, + include_name_in_message=include_name_in_message, + ) + for oai_message_part in oai_message: + for key, value in oai_message_part.items(): + if value is None: + continue + + if isinstance(message, UserMessage) and isinstance(value, list): + typed_message_value = cast(List[ChatCompletionContentPartParam], value) + + assert len(typed_message_value) == len( + message.content + ), "Mismatch in message content and typed message value" + + # We need image properties that are only in the original message + for part, content_part in zip(typed_message_value, message.content, strict=False): + if isinstance(content_part, Image): + # TODO: add detail parameter + num_tokens += calculate_vision_tokens(content_part) + elif isinstance(part, str): + num_tokens += len(encoding.encode(part)) + else: + try: + serialized_part = json.dumps(part) + num_tokens += len(encoding.encode(serialized_part)) + except TypeError: + trace_logger.warning(f"Could not convert {part} to string, skipping.") + else: + if not isinstance(value, str): + try: + value = json.dumps(value) + except TypeError: + trace_logger.warning(f"Could not convert {value} to string, skipping.") + continue + num_tokens += len(encoding.encode(value)) + if key == "name": + num_tokens += tokens_per_name + num_tokens += 3 # every reply is primed with <|start|>assistant<|message|> + + # Tool tokens. + oai_tools = convert_tools(tools) + for tool in oai_tools: + function = tool["function"] + tool_tokens = len(encoding.encode(function["name"])) + if "description" in function: + tool_tokens += len(encoding.encode(function["description"])) + tool_tokens -= 2 + if "parameters" in function: + parameters = function["parameters"] + if "properties" in parameters: + assert isinstance(parameters["properties"], dict) + for propertiesKey in parameters["properties"]: # pyright: ignore + assert isinstance(propertiesKey, str) + tool_tokens += len(encoding.encode(propertiesKey)) + v = parameters["properties"][propertiesKey] # pyright: ignore + for field in v: # pyright: ignore + if field == "type": + tool_tokens += 2 + tool_tokens += len(encoding.encode(v["type"])) # pyright: ignore + elif field == "description": + tool_tokens += 2 + tool_tokens += len(encoding.encode(v["description"])) # pyright: ignore + elif field == "enum": + tool_tokens -= 3 + for o in v["enum"]: # pyright: ignore + tool_tokens += 3 + tool_tokens += len(encoding.encode(o)) # pyright: ignore + else: + trace_logger.warning(f"Not supported field {field}") + tool_tokens += 11 + if len(parameters["properties"]) == 0: # pyright: ignore + tool_tokens -= 2 + num_tokens += tool_tokens + num_tokens += 12 + return num_tokens @dataclass @@ -370,9 +425,11 @@ def __init__( model_capabilities: Optional[ModelCapabilities] = None, # type: ignore model_info: Optional[ModelInfo] = None, add_name_prefixes: bool = False, + include_name_in_message: bool = True, ): self._client = client self._add_name_prefixes = add_name_prefixes + self._include_name_in_message = include_name_in_message if model_capabilities is None and model_info is None: try: self._model_info = _model_info.get_info(create_args["model"]) @@ -417,10 +474,22 @@ def __init__( def create_from_config(cls, config: Dict[str, Any]) -> ChatCompletionClient: return OpenAIChatCompletionClient(**config) + def _rstrip_last_assistant_message(self, messages: Sequence[LLMMessage]) -> Sequence[LLMMessage]: + """ + Remove the last assistant message if it is empty. + """ + # When Claude models last message is AssistantMessage, It could not end with whitespace + if isinstance(messages[-1], AssistantMessage): + if isinstance(messages[-1].content, str): + messages[-1].content = messages[-1].content.rstrip() + + return messages + def _process_create_args( self, messages: Sequence[LLMMessage], tools: Sequence[Tool | ToolSchema], + tool_choice: Tool | Literal["auto", "required", "none"], json_output: Optional[bool | type[BaseModel]], extra_create_args: Mapping[str, Any], ) -> CreateParams: @@ -497,7 +566,50 @@ def _process_create_args( if self.model_info["json_output"] is False and json_output is True: raise ValueError("Model does not support JSON output.") - oai_messages_nested = [to_oai_type(m, prepend_name=self._add_name_prefixes) for m in messages] + if not self.model_info.get("multiple_system_messages", False): + # Some models accept only one system message(or, it will read only the last one) + # So, merge system messages into one (if multiple and continuous) + system_message_content = "" + _messages: List[LLMMessage] = [] + _first_system_message_idx = -1 + _last_system_message_idx = -1 + # Index of the first system message for adding the merged system message at the correct position + for idx, message in enumerate(messages): + if isinstance(message, SystemMessage): + if _first_system_message_idx == -1: + _first_system_message_idx = idx + elif _last_system_message_idx + 1 != idx: + # That case, system message is not continuous + # Merge system messages only contiues system messages + raise ValueError( + "Multiple and Not continuous system messages are not supported if model_info['multiple_system_messages'] is False" + ) + system_message_content += message.content + "\n" + _last_system_message_idx = idx + else: + _messages.append(message) + system_message_content = system_message_content.rstrip() + if system_message_content != "": + system_message = SystemMessage(content=system_message_content) + _messages.insert(_first_system_message_idx, system_message) + messages = _messages + + # in that case, for ad-hoc, we using startswith instead of model_family for code consistency + if create_args.get("model", "unknown").startswith("claude-"): + # When Claude models last message is AssistantMessage, It could not end with whitespace + messages = self._rstrip_last_assistant_message(messages) + + oai_messages_nested = [ + to_oai_type( + m, + prepend_name=self._add_name_prefixes, + model=create_args.get("model", "unknown"), + model_family=self._model_info["family"], + include_name_in_message=self._include_name_in_message, + ) + for m in messages + ] + oai_messages = [item for sublist in oai_messages_nested for item in sublist] if self.model_info["function_calling"] is False and len(tools) > 0: @@ -505,6 +617,29 @@ def _process_create_args( converted_tools = convert_tools(tools) + # Process tool_choice parameter + if isinstance(tool_choice, Tool): + if len(tools) == 0: + raise ValueError("tool_choice specified but no tools provided") + + # Validate that the tool exists in the provided tools + tool_names_available: List[str] = [] + for tool in tools: + if isinstance(tool, Tool): + tool_names_available.append(tool.schema["name"]) + else: + tool_names_available.append(tool["name"]) + + # tool_choice is a single Tool object + tool_name = tool_choice.schema["name"] + if tool_name not in tool_names_available: + raise ValueError(f"tool_choice references '{tool_name}' but it's not in the provided tools") + + if len(converted_tools) > 0: + # Convert to OpenAI format and add to create_args + converted_tool_choice = convert_tool_choice(tool_choice) + create_args["tool_choice"] = converted_tool_choice + return CreateParams( messages=oai_messages, tools=converted_tools, @@ -517,6 +652,7 @@ async def create( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -524,6 +660,7 @@ async def create( create_params = self._process_create_args( messages, tools, + tool_choice, json_output, extra_create_args, ) @@ -555,10 +692,12 @@ async def create( if create_params.response_format is not None: result = cast(ParsedChatCompletion[Any], result) + # Handle the case where OpenAI API might return None for token counts + # even when result.usage is not None usage = RequestUsage( # TODO backup token counting - prompt_tokens=result.usage.prompt_tokens if result.usage is not None else 0, - completion_tokens=(result.usage.completion_tokens if result.usage is not None else 0), + prompt_tokens=getattr(result.usage, "prompt_tokens", 0) if result.usage is not None else 0, + completion_tokens=getattr(result.usage, "completion_tokens", 0) if result.usage is not None else 0, ) logger.info( @@ -567,6 +706,7 @@ async def create( response=result.model_dump(), prompt_tokens=usage.prompt_tokens, completion_tokens=usage.completion_tokens, + tools=create_params.tools, ) ) @@ -665,10 +805,12 @@ async def create_stream( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, max_consecutive_empty_chunk_tolerance: int = 0, + include_usage: Optional[bool] = None, ) -> AsyncGenerator[Union[str, CreateResult], None]: """Create a stream of string chunks from the model ending with a :class:`~autogen_core.models.CreateResult`. @@ -677,7 +819,7 @@ async def create_stream( In streaming, the default behaviour is not return token usage counts. See: `OpenAI API reference for possible args `_. - You can set `extra_create_args={"stream_options": {"include_usage": True}}` + You can set set the `include_usage` flag to True or `extra_create_args={"stream_options": {"include_usage": True}}`. If both the flag and `stream_options` are set, but to different values, an exception will be raised. (if supported by the accessed API) to return a final chunk with usage set to a :class:`~autogen_core.models.RequestUsage` object with prompt and completion token counts, @@ -695,10 +837,22 @@ async def create_stream( create_params = self._process_create_args( messages, tools, + tool_choice, json_output, extra_create_args, ) + if include_usage is not None: + if "stream_options" in create_params.create_args: + stream_options = create_params.create_args["stream_options"] + if "include_usage" in stream_options and stream_options["include_usage"] != include_usage: + raise ValueError( + "include_usage and extra_create_args['stream_options']['include_usage'] are both set, but differ in value." + ) + else: + # If stream options are not present, add them. + create_params.create_args["stream_options"] = {"include_usage": True} + if max_consecutive_empty_chunk_tolerance != 0: warnings.warn( "The 'max_consecutive_empty_chunk_tolerance' parameter is deprecated and will be removed in the future releases. All of empty chunks will be skipped with a warning.", @@ -982,93 +1136,14 @@ def total_usage(self) -> RequestUsage: return self._total_usage def count_tokens(self, messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = []) -> int: - model = self._create_args["model"] - try: - encoding = tiktoken.encoding_for_model(model) - except KeyError: - trace_logger.warning(f"Model {model} not found. Using cl100k_base encoding.") - encoding = tiktoken.get_encoding("cl100k_base") - tokens_per_message = 3 - tokens_per_name = 1 - num_tokens = 0 - - # Message tokens. - for message in messages: - num_tokens += tokens_per_message - oai_message = to_oai_type(message, prepend_name=self._add_name_prefixes) - for oai_message_part in oai_message: - for key, value in oai_message_part.items(): - if value is None: - continue - - if isinstance(message, UserMessage) and isinstance(value, list): - typed_message_value = cast(List[ChatCompletionContentPartParam], value) - - assert len(typed_message_value) == len( - message.content - ), "Mismatch in message content and typed message value" - - # We need image properties that are only in the original message - for part, content_part in zip(typed_message_value, message.content, strict=False): - if isinstance(content_part, Image): - # TODO: add detail parameter - num_tokens += calculate_vision_tokens(content_part) - elif isinstance(part, str): - num_tokens += len(encoding.encode(part)) - else: - try: - serialized_part = json.dumps(part) - num_tokens += len(encoding.encode(serialized_part)) - except TypeError: - trace_logger.warning(f"Could not convert {part} to string, skipping.") - else: - if not isinstance(value, str): - try: - value = json.dumps(value) - except TypeError: - trace_logger.warning(f"Could not convert {value} to string, skipping.") - continue - num_tokens += len(encoding.encode(value)) - if key == "name": - num_tokens += tokens_per_name - num_tokens += 3 # every reply is primed with <|start|>assistant<|message|> - - # Tool tokens. - oai_tools = convert_tools(tools) - for tool in oai_tools: - function = tool["function"] - tool_tokens = len(encoding.encode(function["name"])) - if "description" in function: - tool_tokens += len(encoding.encode(function["description"])) - tool_tokens -= 2 - if "parameters" in function: - parameters = function["parameters"] - if "properties" in parameters: - assert isinstance(parameters["properties"], dict) - for propertiesKey in parameters["properties"]: # pyright: ignore - assert isinstance(propertiesKey, str) - tool_tokens += len(encoding.encode(propertiesKey)) - v = parameters["properties"][propertiesKey] # pyright: ignore - for field in v: # pyright: ignore - if field == "type": - tool_tokens += 2 - tool_tokens += len(encoding.encode(v["type"])) # pyright: ignore - elif field == "description": - tool_tokens += 2 - tool_tokens += len(encoding.encode(v["description"])) # pyright: ignore - elif field == "enum": - tool_tokens -= 3 - for o in v["enum"]: # pyright: ignore - tool_tokens += 3 - tool_tokens += len(encoding.encode(o)) # pyright: ignore - else: - trace_logger.warning(f"Not supported field {field}") - tool_tokens += 11 - if len(parameters["properties"]) == 0: # pyright: ignore - tool_tokens -= 2 - num_tokens += tool_tokens - num_tokens += 12 - return num_tokens + return count_tokens_openai( + messages, + self._create_args["model"], + add_name_prefixes=self._add_name_prefixes, + tools=tools, + model_family=self._model_info["family"], + include_name_in_message=self._include_name_in_message, + ) def remaining_tokens(self, messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = []) -> int: token_limit = _model_info.get_token_limit(self._create_args["model"]) @@ -1161,6 +1236,7 @@ class OpenAIChatCompletionClient(BaseOpenAIChatCompletionClient, Component[OpenA stop (optional, str | List[str]): temperature (optional, float): top_p (optional, float): + parallel_tool_calls (optional, bool): Whether to allow parallel tool calls. When not set, defaults to server behavior. user (optional, str): default_headers (optional, dict[str, str]): Custom headers; useful for authentication or other custom requirements. add_name_prefixes (optional, bool): Whether to prepend the `source` value @@ -1168,6 +1244,9 @@ class OpenAIChatCompletionClient(BaseOpenAIChatCompletionClient, Component[OpenA "this is content" becomes "Reviewer said: this is content." This can be useful for models that do not support the `name` field in message. Defaults to False. + include_name_in_message (optional, bool): Whether to include the `name` field + in user message parameters sent to the OpenAI API. Defaults to True. Set to False + for model providers that don't support the `name` field (e.g., Groq). stream_options (optional, dict): Additional options for streaming. Currently only `include_usage` is supported. Examples: @@ -1367,6 +1446,10 @@ def __init__(self, **kwargs: Unpack[OpenAIClientConfiguration]): if "add_name_prefixes" in kwargs: add_name_prefixes = kwargs["add_name_prefixes"] + include_name_in_message: bool = True + if "include_name_in_message" in kwargs: + include_name_in_message = kwargs["include_name_in_message"] + # Special handling for Gemini model. assert "model" in copied_args and isinstance(copied_args["model"], str) if copied_args["model"].startswith("gemini-"): @@ -1374,6 +1457,16 @@ def __init__(self, **kwargs: Unpack[OpenAIClientConfiguration]): copied_args["base_url"] = _model_info.GEMINI_OPENAI_BASE_URL if "api_key" not in copied_args and "GEMINI_API_KEY" in os.environ: copied_args["api_key"] = os.environ["GEMINI_API_KEY"] + if copied_args["model"].startswith("claude-"): + if "base_url" not in copied_args: + copied_args["base_url"] = _model_info.ANTHROPIC_OPENAI_BASE_URL + if "api_key" not in copied_args and "ANTHROPIC_API_KEY" in os.environ: + copied_args["api_key"] = os.environ["ANTHROPIC_API_KEY"] + if copied_args["model"].startswith("Llama-"): + if "base_url" not in copied_args: + copied_args["base_url"] = _model_info.LLAMA_API_BASE_URL + if "api_key" not in copied_args and "LLAMA_API_KEY" in os.environ: + copied_args["api_key"] = os.environ["LLAMA_API_KEY"] client = _openai_client_from_config(copied_args) create_args = _create_args_from_config(copied_args) @@ -1384,6 +1477,7 @@ def __init__(self, **kwargs: Unpack[OpenAIClientConfiguration]): model_capabilities=model_capabilities, model_info=model_info, add_name_prefixes=add_name_prefixes, + include_name_in_message=include_name_in_message, ) def __getstate__(self) -> Dict[str, Any]: @@ -1483,8 +1577,18 @@ class AzureOpenAIChatCompletionClient( stop (optional, str | List[str]): temperature (optional, float): top_p (optional, float): + parallel_tool_calls (optional, bool): Whether to allow parallel tool calls. When not set, defaults to server behavior. user (optional, str): default_headers (optional, dict[str, str]): Custom headers; useful for authentication or other custom requirements. + add_name_prefixes (optional, bool): Whether to prepend the `source` value + to each :class:`~autogen_core.models.UserMessage` content. E.g., + "this is content" becomes "Reviewer said: this is content." + This can be useful for models that do not support the `name` field in + message. Defaults to False. + include_name_in_message (optional, bool): Whether to include the `name` field + in user message parameters sent to the OpenAI API. Defaults to True. Set to False + for model providers that don't support the `name` field (e.g., Groq). + stream_options (optional, dict): Additional options for streaming. Currently only `include_usage` is supported. To use the client, you need to provide your deployment name, Azure Cognitive Services endpoint, and api version. @@ -1548,6 +1652,10 @@ class AzureOpenAIChatCompletionClient( Right now only `DefaultAzureCredential` is supported with no additional args passed to it. + .. note:: + + The Azure OpenAI client by default sets the User-Agent header to `autogen-python/{version}`. To override this, you can set the variable `autogen_ext.models.openai.AZURE_OPENAI_USER_AGENT` environment variable to an empty string. + See `here `_ for how to use the Azure client directly or for more info. """ @@ -1572,6 +1680,10 @@ def __init__(self, **kwargs: Unpack[AzureOpenAIClientConfiguration]): if "add_name_prefixes" in kwargs: add_name_prefixes = kwargs["add_name_prefixes"] + include_name_in_message: bool = True + if "include_name_in_message" in kwargs: + include_name_in_message = kwargs["include_name_in_message"] + client = _azure_openai_client_from_config(copied_args) create_args = _create_args_from_config(copied_args) self._raw_config: Dict[str, Any] = copied_args @@ -1581,6 +1693,7 @@ def __init__(self, **kwargs: Unpack[AzureOpenAIClientConfiguration]): model_capabilities=model_capabilities, model_info=model_info, add_name_prefixes=add_name_prefixes, + include_name_in_message=include_name_in_message, ) def __getstate__(self) -> Dict[str, Any]: diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_transformation/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_transformation/__init__.py new file mode 100644 index 000000000000..dc21b9c10815 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_transformation/__init__.py @@ -0,0 +1,27 @@ +from .registry import ( + MESSAGE_TRANSFORMERS, + build_conditional_transformer_func, + build_transformer_func, + get_transformer, + register_transformer, +) +from .types import ( + LLMMessageContent, + MessageParam, + TransformerFunc, + TransformerMap, + TrasformerReturnType, +) + +__all__ = [ + "register_transformer", + "get_transformer", + "build_transformer_func", + "build_conditional_transformer_func", + "MESSAGE_TRANSFORMERS", + "TransformerMap", + "TransformerFunc", + "MessageParam", + "LLMMessageContent", + "TrasformerReturnType", +] diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_transformation/registry.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_transformation/registry.py new file mode 100644 index 000000000000..bc603f4a55ec --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_transformation/registry.py @@ -0,0 +1,130 @@ +from collections import defaultdict +from typing import Any, Callable, Dict, List, get_args + +from autogen_core.models import LLMMessage, ModelFamily + +from .types import ( + TransformerFunc, + TransformerMap, +) + +# Global registry of model family → message transformer map +# Each model family (e.g. "gpt-4o", "gemini-1.5-flash") maps to a dict of LLMMessage type → transformer function +MESSAGE_TRANSFORMERS: Dict[str, Dict[str, TransformerMap]] = defaultdict(dict) + + +def build_transformer_func( + funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]], message_param_func: Callable[..., Any] +) -> TransformerFunc: + """ + Combines multiple transformer functions into a single transformer. + + Each `func` must accept a message and a context dict, and return a partial dict + of keyword arguments. These are merged and passed to `message_param_func`. + + This structure allows flexible transformation pipelines and future extensibility + (e.g., prepend name, insert metadata, etc). + + message_param_func: A model-specific constructor (e.g. ChatCompletionMessageParam). + Signature is intentionally open: Callable[..., Any]. + """ + + def transformer_func(message: LLMMessage, context: Any) -> Any: + kwargs: Dict[str, Any] = {} + for func in funcs: + kwargs.update(func(message, context)) + return [message_param_func(**kwargs)] + + return transformer_func + + +def build_conditional_transformer_func( + funcs_map: Dict[str, List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]]], + message_param_func_map: Dict[str, Callable[..., Any]], + condition_func: Callable[[LLMMessage, Dict[str, Any]], str], +) -> TransformerFunc: + """ + Combines multiple transformer functions into a single transformer, with a conditional constructor. + + Each `func` must accept a message and a context dict, and return a partial dict + of keyword arguments. These are merged and passed to the constructor selected by `condition_func`. + + This structure allows flexible transformation pipelines and future extensibility + (e.g., prepend name, insert metadata, etc). + + message_param_func_map: A mapping of condition → constructor function. + condition_func: A function that returns the condition for selecting the constructor. + """ + + def transformer(message: LLMMessage, context: Dict[str, Any]) -> Any: + condition = condition_func(message, context) + message_param_func = message_param_func_map[condition] + kwargs: Dict[str, Any] = {} + for func in funcs_map[condition]: + kwargs.update(func(message, context)) + if kwargs.get("pass_message", False): + return [] + return [message_param_func(**kwargs)] + + return transformer + + +def register_transformer(api: str, model_family: str, transformer_map: TransformerMap) -> None: + """ + Registers a transformer map for a given model family. + + Example: + + .. code-block:: python + + register_transformer( + "gpt-4o", + { + UserMessage: user_message_to_oai, + SystemMessage: system_message_to_oai, + }, + ) + """ + MESSAGE_TRANSFORMERS[api][model_family] = transformer_map + + +def _find_model_family(api: str, model: str) -> str: + """ + Finds the best matching model family for the given model. + Search via prefix matching (e.g. "gpt-4o" → "gpt-4o-1.0"). + """ + len_family = 0 + family = ModelFamily.UNKNOWN + for _family in MESSAGE_TRANSFORMERS[api].keys(): + if model.startswith(_family): + if len(_family) > len_family: + family = _family + len_family = len(_family) + return family + + +def get_transformer(api: str, model: str, model_family: str) -> TransformerMap: + """ + Returns the registered transformer map for the given model family. + + This is a thin wrapper around `MESSAGE_TRANSFORMERS.get(...)`, but serves as + an abstraction layer to allow future enhancements such as: + + - Providing fallback transformers for unknown model families + - Injecting mock transformers during testing + - Adding logging, metrics, or versioning later + + Keeping this as a function (instead of direct dict access) improves long-term flexibility. + """ + + if model_family not in set(get_args(ModelFamily.ANY)) or model_family == ModelFamily.UNKNOWN: + # fallback to finding the best matching model family + model_family = _find_model_family(api, model) + + transformer = MESSAGE_TRANSFORMERS.get(api, {}).get(model_family, {}) + + if not transformer: + # Just in case, we should never reach here + raise ValueError(f"No transformer found for model family '{model_family}'") + + return transformer diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_transformation/types.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_transformation/types.py new file mode 100644 index 000000000000..9cfb28e040cc --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_transformation/types.py @@ -0,0 +1,22 @@ +from typing import Any, Callable, Dict, List, Sequence, Type, Union + +from autogen_core import FunctionCall, Image +from autogen_core.models import LLMMessage +from autogen_core.models._types import FunctionExecutionResult +from openai.types.chat import ChatCompletionMessageParam + +MessageParam = Union[ChatCompletionMessageParam] # If that transformation move to global, add other message params here +TrasformerReturnType = Sequence[MessageParam] +TransformerFunc = Callable[[LLMMessage, Dict[str, Any]], TrasformerReturnType] +TransformerMap = Dict[Type[LLMMessage], TransformerFunc] + +LLMMessageContent = Union[ + # SystemMessage.content + str, + # UserMessage.content + List[Union[str, Image]], + # AssistantMessage.content + List[FunctionCall], + # FunctionExecutionResultMessage.content + List[FunctionExecutionResult], +] diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_utils.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_utils.py new file mode 100644 index 000000000000..8c1df22961d7 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_utils.py @@ -0,0 +1,14 @@ +import re + + +def assert_valid_name(name: str) -> str: + """ + Ensure that configured names are valid, raises ValueError if not. + + For munging LLM responses use _normalize_name to ensure LLM specified names don't break the API. + """ + if not re.match(r"^[a-zA-Z0-9_-]+$", name): + raise ValueError(f"Invalid name: {name}. Only letters, numbers, '_' and '-' are allowed.") + if len(name) > 64: + raise ValueError(f"Invalid name: {name}. Name must be less than 64 characters.") + return name diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/config/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/config/__init__.py index a125105255ca..02b1e3a80ff6 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/config/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/config/__init__.py @@ -49,6 +49,7 @@ class CreateArguments(TypedDict, total=False): top_p: Optional[float] user: str stream_options: Optional[StreamOptions] + parallel_tool_calls: Optional[bool] AsyncAzureADTokenProvider = Callable[[], Union[str, Awaitable[str]]] @@ -63,6 +64,8 @@ class BaseOpenAIClientConfiguration(CreateArguments, total=False): model_info: ModelInfo add_name_prefixes: bool """What functionality the model supports, determined by default from model name but is overriden if value passed.""" + include_name_in_message: bool + """Whether to include the 'name' field in user message parameters. Defaults to True. Set to False for providers that don't support the 'name' field.""" default_headers: Dict[str, str] | None @@ -95,6 +98,7 @@ class CreateArgumentsConfigModel(BaseModel): top_p: float | None = None user: str | None = None stream_options: StreamOptions | None = None + parallel_tool_calls: bool | None = None class BaseOpenAIClientConfigurationConfigModel(CreateArgumentsConfigModel): @@ -105,6 +109,7 @@ class BaseOpenAIClientConfigurationConfigModel(CreateArgumentsConfigModel): model_capabilities: ModelCapabilities | None = None # type: ignore model_info: ModelInfo | None = None add_name_prefixes: bool | None = None + include_name_in_message: bool | None = None default_headers: Dict[str, str] | None = None diff --git a/python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py b/python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py index 6e2b03beb2c7..8aaa25a9d098 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py @@ -2,7 +2,7 @@ import logging import warnings -from typing import Any, AsyncGenerator, Dict, List, Mapping, Optional, Sequence, Union +from typing import Any, AsyncGenerator, Dict, List, Literal, Mapping, Optional, Sequence, Union from autogen_core import EVENT_LOGGER_NAME, CancellationToken, Component from autogen_core.models import ( @@ -162,11 +162,16 @@ async def create( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, ) -> CreateResult: """Return the next completion from the list.""" + # Warn if tool_choice is specified since it's ignored in replay mode + if tool_choice != "auto": + logger.warning("tool_choice parameter specified but is ignored in replay mode") + if self._current_index >= len(self.chat_completions): raise ValueError("No more mock responses available") @@ -201,11 +206,16 @@ async def create_stream( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, ) -> AsyncGenerator[Union[str, CreateResult], None]: """Return the next completion as a stream.""" + # Warn if tool_choice is specified since it's ignored in replay mode + if tool_choice != "auto": + logger.warning("tool_choice parameter specified but is ignored in replay mode") + if self._current_index >= len(self.chat_completions): raise ValueError("No more mock responses available") diff --git a/python/packages/autogen-ext/src/autogen_ext/models/semantic_kernel/_sk_chat_completion_adapter.py b/python/packages/autogen-ext/src/autogen_ext/models/semantic_kernel/_sk_chat_completion_adapter.py index 78f0aa5a24de..b9267057cd49 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/semantic_kernel/_sk_chat_completion_adapter.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/semantic_kernel/_sk_chat_completion_adapter.py @@ -1,7 +1,7 @@ import json import logging import warnings -from typing import Any, Literal, Mapping, Optional, Sequence +from typing import Any, Literal, Mapping, Optional, Sequence, Union from autogen_core import EVENT_LOGGER_NAME, FunctionCall from autogen_core._cancellation_token import CancellationToken @@ -29,9 +29,9 @@ ) from semantic_kernel.functions.kernel_plugin import KernelPlugin from semantic_kernel.kernel import Kernel -from typing_extensions import AsyncGenerator, Union +from typing_extensions import AsyncGenerator -from autogen_ext.tools.semantic_kernel import KernelFunctionFromTool +from autogen_ext.tools.semantic_kernel import KernelFunctionFromTool, KernelFunctionFromToolSchema from .._utils.parse_r1_content import parse_r1_content @@ -396,6 +396,9 @@ def _sync_tools_with_kernel(self, kernel: Kernel, tools: Sequence[Tool | ToolSch # Convert Tool to KernelFunction using KernelFunctionFromTool kernel_function = KernelFunctionFromTool(tool) # type: ignore self._tools_plugin.functions[tool.schema["name"]] = kernel_function + else: + kernel_function = KernelFunctionFromToolSchema(tool) # type: ignore + self._tools_plugin.functions[tool.get("name")] = kernel_function # type: ignore kernel.add_plugin(self._tools_plugin) @@ -439,6 +442,7 @@ async def create( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -470,6 +474,13 @@ async def create( if isinstance(json_output, type) and issubclass(json_output, BaseModel): raise ValueError("structured output is not currently supported in SKChatCompletionAdapter") + # Handle tool_choice parameter + if tool_choice != "auto": + warnings.warn( + "tool_choice parameter is specified but may not be fully supported by SKChatCompletionAdapter.", + stacklevel=2, + ) + kernel = self._get_kernel(extra_create_args) chat_history = self._convert_to_chat_history(messages) @@ -550,6 +561,7 @@ async def create_stream( messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = [], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", json_output: Optional[bool | type[BaseModel]] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, @@ -582,6 +594,13 @@ async def create_stream( if isinstance(json_output, type) and issubclass(json_output, BaseModel): raise ValueError("structured output is not currently supported in SKChatCompletionAdapter") + # Handle tool_choice parameter + if tool_choice != "auto": + warnings.warn( + "tool_choice parameter is specified but may not be fully supported by SKChatCompletionAdapter.", + stacklevel=2, + ) + kernel = self._get_kernel(extra_create_args) chat_history = self._convert_to_chat_history(messages) user_settings = self._get_prompt_settings(extra_create_args) diff --git a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime.py b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime.py index 504838740283..6a3963586e18 100644 --- a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime.py +++ b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime.py @@ -251,6 +251,7 @@ def __init__( self._subscription_manager = SubscriptionManager() self._serialization_registry = SerializationRegistry() self._extra_grpc_config = extra_grpc_config or [] + self._agent_instance_types: Dict[str, Type[Agent]] = {} if payload_serialization_format not in {JSON_DATA_CONTENT_TYPE, PROTOBUF_DATA_CONTENT_TYPE}: raise ValueError(f"Unsupported payload serialization format: {payload_serialization_format}") @@ -701,6 +702,14 @@ async def send_message(agent: Agent, message_context: MessageContext) -> Any: except BaseException as e: logger.error("Error handling event", exc_info=e) + async def _register_agent_type(self, agent_type: str) -> None: + if self._host_connection is None: + raise RuntimeError("Host connection is not set.") + message = agent_worker_pb2.RegisterAgentTypeRequest(type=agent_type) + _response: agent_worker_pb2.RegisterAgentTypeResponse = await self._host_connection.stub.RegisterAgent( + message, metadata=self._host_connection.metadata + ) + async def register_factory( self, type: str | AgentType, @@ -729,14 +738,38 @@ async def factory_wrapper() -> T: return agent_instance self._agent_factories[type.type] = factory_wrapper - # Send the registration request message to the host. - message = agent_worker_pb2.RegisterAgentTypeRequest(type=type.type) - _response: agent_worker_pb2.RegisterAgentTypeResponse = await self._host_connection.stub.RegisterAgent( - message, metadata=self._host_connection.metadata - ) + await self._register_agent_type(type.type) + return type + async def register_agent_instance( + self, + agent_instance: Agent, + agent_id: AgentId, + ) -> AgentId: + def agent_factory() -> Agent: + raise RuntimeError( + "Agent factory was invoked for an agent instance that was not registered. This is likely due to the agent type being incorrectly subscribed to a topic. If this exception occurs when publishing a message to the DefaultTopicId, then it is likely that `skip_class_subscriptions` needs to be turned off when registering the agent." + ) + + if agent_id in self._instantiated_agents: + raise ValueError(f"Agent with id {agent_id} already exists.") + + if agent_id.type not in self._agent_factories: + self._agent_factories[agent_id.type] = agent_factory + await self._register_agent_type(agent_id.type) + self._agent_instance_types[agent_id.type] = type_func_alias(agent_instance) + else: + if self._agent_factories[agent_id.type].__code__ != agent_factory.__code__: + raise ValueError("Agent factories and agent instances cannot be registered to the same type.") + if self._agent_instance_types[agent_id.type] != type_func_alias(agent_instance): + raise ValueError("Agent instances must be the same object type.") + + await agent_instance.bind_id_and_runtime(id=agent_id, runtime=self) + self._instantiated_agents[agent_id] = agent_instance + return agent_id + async def _invoke_agent_factory( self, agent_factory: Callable[[], T | Awaitable[T]] | Callable[[AgentRuntime, AgentId], T | Awaitable[T]], @@ -757,7 +790,7 @@ async def _invoke_agent_factory( raise ValueError("Agent factory must take 0 or 2 arguments.") if inspect.isawaitable(agent): - return cast(T, await agent) + agent = cast(T, await agent) return agent diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one.py b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one.py index ec2d5192ab74..5acc709550a9 100644 --- a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one.py +++ b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one.py @@ -1,7 +1,7 @@ import warnings from typing import Awaitable, Callable, List, Optional, Union -from autogen_agentchat.agents import CodeExecutorAgent, UserProxyAgent +from autogen_agentchat.agents import ApprovalFuncType, CodeExecutorAgent, UserProxyAgent from autogen_agentchat.base import ChatAgent from autogen_agentchat.teams import MagenticOneGroupChat from autogen_core import CancellationToken @@ -14,11 +14,61 @@ from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor from autogen_ext.models.openai._openai_client import BaseOpenAIChatCompletionClient +# Docker imports for default code executor +try: + import docker + from docker.errors import DockerException + + from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor + + _docker_available = True +except ImportError: + docker = None # type: ignore + DockerException = Exception # type: ignore + DockerCommandLineCodeExecutor = None # type: ignore + _docker_available = False + SyncInputFunc = Callable[[str], str] AsyncInputFunc = Callable[[str, Optional[CancellationToken]], Awaitable[str]] InputFuncType = Union[SyncInputFunc, AsyncInputFunc] +def _is_docker_available() -> bool: + """Check if Docker is available and running.""" + if not _docker_available: + return False + + try: + if docker is not None: + client = docker.from_env() + client.ping() # type: ignore + return True + except DockerException: + return False + + return False + + +def _create_default_code_executor() -> CodeExecutor: + """Create the default code executor, preferring Docker if available.""" + if _is_docker_available() and DockerCommandLineCodeExecutor is not None: + try: + return DockerCommandLineCodeExecutor() + except Exception: + # Fallback to local if Docker fails to initialize + pass + + # Issue warning and use local executor if Docker is not available + warnings.warn( + "Docker is not available or not running. Using LocalCommandLineCodeExecutor instead of the recommended DockerCommandLineCodeExecutor. " + "For security, it is recommended to install Docker and ensure it's running before using MagenticOne. " + "To install Docker, visit: https://docs.docker.com/get-docker/", + UserWarning, + stacklevel=3, + ) + return LocalCommandLineCodeExecutor() + + class MagenticOne(MagenticOneGroupChat): """ MagenticOne is a specialized group chat class that integrates various agents @@ -35,6 +85,9 @@ class MagenticOne(MagenticOneGroupChat): Args: client (ChatCompletionClient): The client used for model interactions. hil_mode (bool): Optional; If set to True, adds the UserProxyAgent to the list of agents. + input_func (InputFuncType | None): Optional; Function to use for user input in human-in-the-loop mode. + code_executor (CodeExecutor | None): Optional; Code executor to use. If None, will use Docker if available, otherwise local executor. + approval_func (ApprovalFuncType | None): Optional; Function to approve code execution before running. If None, code will execute without approval. .. warning:: Using Magentic-One involves interacting with a digital world designed for humans, which carries inherent risks. To minimize these risks, consider the following precautions: @@ -77,7 +130,7 @@ class MagenticOne(MagenticOneGroupChat): async def example_usage(): client = OpenAIChatCompletionClient(model="gpt-4o") - m1 = MagenticOne(client=client) + m1 = MagenticOne(client=client) # Uses DockerCommandLineCodeExecutor by default task = "Write a Python script to fetch data from an API." result = await Console(m1.run_stream(task=task)) print(result) @@ -89,25 +142,89 @@ async def example_usage(): .. code-block:: python - # Enable human-in-the-loop mode + # Enable human-in-the-loop mode with explicit Docker executor and code approval import asyncio from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.teams.magentic_one import MagenticOne + from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor from autogen_agentchat.ui import Console + from autogen_agentchat.agents import ApprovalRequest, ApprovalResponse + + + def user_input_func(prompt: str) -> str: + \"\"\"Custom input function for user interaction.\"\"\" + return input(prompt) + + + def approval_func(request: ApprovalRequest) -> ApprovalResponse: + \"\"\"Simple approval function that requests user input.\"\"\" + print(f\"Code to execute:\\n{request.code}\") + user_input = input("Do you approve this code execution? (y/n): ").strip().lower() + if user_input == 'y': + return ApprovalResponse(approved=True, reason=\"User approved the code execution\") + else: + return ApprovalResponse(approved=False, reason=\"User denied the code execution\") async def example_usage_hil(): client = OpenAIChatCompletionClient(model="gpt-4o") - # to enable human-in-the-loop mode, set hil_mode=True - m1 = MagenticOne(client=client, hil_mode=True) - task = "Write a Python script to fetch data from an API." - result = await Console(m1.run_stream(task=task)) - print(result) + # Explicitly specify Docker code executor for better security + async with DockerCommandLineCodeExecutor() as code_executor: + m1 = MagenticOne( + client=client, + hil_mode=True, + input_func=user_input_func, + code_executor=code_executor, + approval_func=approval_func + ) + task = "Write a Python script to fetch data from an API." + result = await Console(m1.run_stream(task=task)) + print(result) if __name__ == "__main__": asyncio.run(example_usage_hil()) + + .. code-block:: python + + # Enable code execution approval without human-in-the-loop mode + import asyncio + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_ext.teams.magentic_one import MagenticOne + from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor + from autogen_agentchat.ui import Console + from autogen_agentchat.agents import ApprovalRequest, ApprovalResponse + + + def approval_func(request: ApprovalRequest) -> ApprovalResponse: + \"\"\"Simple approval function that requests user input.\"\"\" + print(f\"Code to execute:\\n{request.code}\") + user_input = input("Do you approve this code execution? (y/n): ").strip().lower() + if user_input == 'y': + return ApprovalResponse(approved=True, reason=\"User approved the code execution\") + else: + return ApprovalResponse(approved=False, reason=\"User denied the code execution\") + + + async def example_usage_with_approval(): + client = OpenAIChatCompletionClient(model="gpt-4o") + # Use approval_func for code approval only (hil_mode=False) + async with DockerCommandLineCodeExecutor() as code_executor: + m1 = MagenticOne( + client=client, + hil_mode=False, # No human-in-the-loop for general conversation + code_executor=code_executor, + approval_func=approval_func # But still ask for code execution approval + ) + task = "Write a Python script to fetch data from an API." + result = await Console(m1.run_stream(task=task)) + print(result) + + + if __name__ == "__main__": + asyncio.run(example_usage_with_approval()) + References: .. code-block:: bibtex @@ -128,22 +245,24 @@ def __init__( hil_mode: bool = False, input_func: InputFuncType | None = None, code_executor: CodeExecutor | None = None, + approval_func: ApprovalFuncType | None = None, ): self.client = client self._validate_client_capabilities(client) if code_executor is None: warnings.warn( - "Instantiating MagenticOne without a code_executor is deprecated. Provide a code_executor to clear this warning (e.g., code_executor=LocalCommandLineCodeExecutor() ).", + "Instantiating MagenticOne without a code_executor is deprecated. Provide a code_executor to clear this warning (e.g., code_executor=DockerCommandLineCodeExecutor() ).", DeprecationWarning, stacklevel=2, ) - code_executor = LocalCommandLineCodeExecutor() + code_executor = _create_default_code_executor() fs = FileSurfer("FileSurfer", model_client=client) ws = MultimodalWebSurfer("WebSurfer", model_client=client) coder = MagenticOneCoderAgent("Coder", model_client=client) - executor = CodeExecutorAgent("ComputerTerminal", code_executor=code_executor) + + executor = CodeExecutorAgent("ComputerTerminal", code_executor=code_executor, approval_func=approval_func) agents: List[ChatAgent] = [fs, ws, coder, executor] if hil_mode: diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/azure/__init__.py b/python/packages/autogen-ext/src/autogen_ext/tools/azure/__init__.py new file mode 100644 index 000000000000..12851185e60c --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/tools/azure/__init__.py @@ -0,0 +1,19 @@ +from ._ai_search import ( + AzureAISearchTool, + BaseAzureAISearchTool, + SearchQuery, + SearchResult, + SearchResults, + VectorizableTextQuery, +) +from ._config import AzureAISearchConfig + +__all__ = [ + "AzureAISearchTool", + "BaseAzureAISearchTool", + "SearchQuery", + "SearchResult", + "SearchResults", + "AzureAISearchConfig", + "VectorizableTextQuery", +] diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/azure/_ai_search.py b/python/packages/autogen-ext/src/autogen_ext/tools/azure/_ai_search.py new file mode 100644 index 000000000000..eb77f78df79c --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/tools/azure/_ai_search.py @@ -0,0 +1,1137 @@ +from __future__ import annotations + +import asyncio +import logging +import time +from abc import ABC, abstractmethod +from contextvars import ContextVar +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + Optional, + Protocol, + Union, +) + +from autogen_core import CancellationToken, Component +from autogen_core.tools import BaseTool, ToolSchema +from pydantic import BaseModel, Field + +from azure.core.credentials import AzureKeyCredential +from azure.core.credentials_async import AsyncTokenCredential +from azure.core.exceptions import HttpResponseError, ResourceNotFoundError +from azure.search.documents.aio import SearchClient + +from ._config import ( + DEFAULT_API_VERSION, + AzureAISearchConfig, +) + +SearchDocument = Dict[str, Any] +MetadataDict = Dict[str, Any] +ContentDict = Dict[str, Any] + +if TYPE_CHECKING: + from azure.search.documents.aio import AsyncSearchItemPaged + + SearchResultsIterable = AsyncSearchItemPaged[SearchDocument] +else: + SearchResultsIterable = Any + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from azure.search.documents.models import ( + VectorizableTextQuery, + VectorizedQuery, + VectorQuery, + ) + +try: + from azure.search.documents.models import VectorizableTextQuery, VectorizedQuery, VectorQuery + + has_azure_search = True +except ImportError: + has_azure_search = False + logger.error( + "The 'azure-search-documents' package is required for this tool but was not found. " + "Please install it with: uv add install azure-search-documents" + ) + + +if TYPE_CHECKING: + from typing import Protocol + + class SearchClientProtocol(Protocol): + async def search(self, **kwargs: Any) -> SearchResultsIterable: ... + async def close(self) -> None: ... +else: + SearchClientProtocol = Any + +__all__ = [ + "AzureAISearchTool", + "BaseAzureAISearchTool", + "SearchQuery", + "SearchResults", + "SearchResult", + "VectorizableTextQuery", + "VectorizedQuery", + "VectorQuery", +] +logger = logging.getLogger(__name__) + + +class SearchQuery(BaseModel): + """Search query parameters. + + This simplified interface only requires a search query string. + All other parameters (top, filters, vector fields, etc.) are specified during tool creation + rather than at query time, making it easier for language models to generate structured output. + + Args: + query (str): The search query text. + """ + + query: str = Field(description="Search query text") + + +class SearchResult(BaseModel): + """Search result. + + Args: + score (float): The search score. + content (ContentDict): The document content. + metadata (MetadataDict): Additional metadata about the document. + """ + + score: float = Field(description="The search score") + content: ContentDict = Field(description="The document content") + metadata: MetadataDict = Field(description="Additional metadata about the document") + + +class SearchResults(BaseModel): + """Container for search results. + + Args: + results (List[SearchResult]): List of search results. + """ + + results: List[SearchResult] = Field(description="List of search results") + + +class EmbeddingProvider(Protocol): + """Protocol defining the interface for embedding generation.""" + + async def _get_embedding(self, query: str) -> List[float]: + """Generate embedding vector for the query text.""" + ... + + +class EmbeddingProviderMixin: + """Mixin class providing embedding generation functionality.""" + + search_config: AzureAISearchConfig + + async def _get_embedding(self, query: str) -> List[float]: + """Generate embedding vector for the query text.""" + if not hasattr(self, "search_config"): + raise ValueError("Host class must have a search_config attribute") + + search_config = self.search_config + embedding_provider = getattr(search_config, "embedding_provider", None) + embedding_model = getattr(search_config, "embedding_model", None) + + if not embedding_provider or not embedding_model: + raise ValueError( + "Client-side embedding is not configured. `embedding_provider` and `embedding_model` must be set." + ) from None + + if embedding_provider.lower() == "azure_openai": + try: + from openai import AsyncAzureOpenAI + + from azure.identity import DefaultAzureCredential + except ImportError: + raise ImportError( + "Azure OpenAI SDK is required for client-side embedding generation. " + "Please install it with: uv add openai azure-identity" + ) from None + + api_key = getattr(search_config, "openai_api_key", None) + api_version = getattr(search_config, "openai_api_version", "2023-11-01") + endpoint = getattr(search_config, "openai_endpoint", None) + + if not endpoint: + raise ValueError( + "Azure OpenAI endpoint (`openai_endpoint`) must be provided for client-side Azure OpenAI embeddings." + ) from None + + if api_key: + azure_client = AsyncAzureOpenAI(api_key=api_key, api_version=api_version, azure_endpoint=endpoint) + else: + + def get_token() -> str: + credential = DefaultAzureCredential() + token = credential.get_token("https://cognitiveservices.azure.com/.default") + if not token or not token.token: + raise ValueError("Failed to acquire token using DefaultAzureCredential for Azure OpenAI.") + return token.token + + azure_client = AsyncAzureOpenAI( + azure_ad_token_provider=get_token, api_version=api_version, azure_endpoint=endpoint + ) + + try: + response = await azure_client.embeddings.create(model=embedding_model, input=query) + return response.data[0].embedding + except Exception as e: + raise ValueError(f"Failed to generate embeddings with Azure OpenAI: {str(e)}") from e + + elif embedding_provider.lower() == "openai": + try: + from openai import AsyncOpenAI + except ImportError: + raise ImportError( + "OpenAI SDK is required for client-side embedding generation. " + "Please install it with: uv add openai" + ) from None + + api_key = getattr(search_config, "openai_api_key", None) + openai_client = AsyncOpenAI(api_key=api_key) + + try: + response = await openai_client.embeddings.create(model=embedding_model, input=query) + return response.data[0].embedding + except Exception as e: + raise ValueError(f"Failed to generate embeddings with OpenAI: {str(e)}") from e + else: + raise ValueError( + f"Unsupported client-side embedding provider: {embedding_provider}. " + "Currently supported providers are 'azure_openai' and 'openai'." + ) + + +class BaseAzureAISearchTool( + BaseTool[SearchQuery, SearchResults], Component[AzureAISearchConfig], EmbeddingProvider, ABC +): + """Abstract base class for Azure AI Search tools. + + This class defines the common interface and functionality for all Azure AI Search tools. + It handles configuration management, client initialization, and the abstract methods + that subclasses must implement. + + Attributes: + search_config: Configuration parameters for the search service. + + Note: + This is an abstract base class and should not be instantiated directly. + Use concrete implementations or the factory methods in AzureAISearchTool. + """ + + component_config_schema = AzureAISearchConfig + component_provider_override = "autogen_ext.tools.azure.BaseAzureAISearchTool" + + def __init__( + self, + name: str, + endpoint: str, + index_name: str, + credential: Union[AzureKeyCredential, AsyncTokenCredential, Dict[str, str]], + description: Optional[str] = None, + api_version: str = DEFAULT_API_VERSION, + query_type: Literal["simple", "full", "semantic", "vector"] = "simple", + search_fields: Optional[List[str]] = None, + select_fields: Optional[List[str]] = None, + vector_fields: Optional[List[str]] = None, + top: Optional[int] = None, + filter: Optional[str] = None, + semantic_config_name: Optional[str] = None, + enable_caching: bool = False, + cache_ttl_seconds: int = 300, + embedding_provider: Optional[str] = None, + embedding_model: Optional[str] = None, + openai_api_key: Optional[str] = None, + openai_api_version: Optional[str] = None, + openai_endpoint: Optional[str] = None, + ): + """Initialize the Azure AI Search tool. + + Args: + name (str): The name of this tool instance + endpoint (str): The full URL of your Azure AI Search service + index_name (str): Name of the search index to query + credential (Union[AzureKeyCredential, TokenCredential, Dict[str, str]]): Azure credential for authentication + description (Optional[str]): Optional description explaining the tool's purpose + api_version (Optional[str]): Azure AI Search API version to use + query_type (Literal["simple", "full", "semantic", "vector"]): Type of search to perform + search_fields (Optional[List[str]]): Fields to search within documents + select_fields (Optional[List[str]]): Fields to return in search results + vector_fields (Optional[List[str]]): Fields to use for vector search + top (Optional[int]): Maximum number of results to return + filter (Optional[str]): OData filter expression to refine search results + semantic_config_name (Optional[str]): Semantic configuration name for enhanced results + enable_caching (bool): Whether to cache search results + cache_ttl_seconds (int): How long to cache results in seconds + embedding_provider (Optional[str]): Name of embedding provider for client-side embeddings + embedding_model (Optional[str]): Model name for client-side embeddings + openai_api_key (Optional[str]): API key for OpenAI/Azure OpenAI embeddings + openai_api_version (Optional[str]): API version for Azure OpenAI embeddings + openai_endpoint (Optional[str]): Endpoint URL for Azure OpenAI embeddings + """ + if not has_azure_search: + raise ImportError( + "Azure Search SDK is required but not installed. " + "Please install it with: pip install azure-search-documents>=11.4.0" + ) + + if description is None: + description = ( + f"Search for information in the {index_name} index using Azure AI Search. " + f"Supports full-text search with optional filters and semantic capabilities." + ) + + super().__init__( + args_type=SearchQuery, + return_type=SearchResults, + name=name, + description=description, + ) + + processed_credential = self._process_credential(credential) + + self.search_config: AzureAISearchConfig = AzureAISearchConfig( + name=name, + description=description, + endpoint=endpoint, + index_name=index_name, + credential=processed_credential, + api_version=api_version, + query_type=query_type, + search_fields=search_fields, + select_fields=select_fields, + vector_fields=vector_fields, + top=top, + filter=filter, + semantic_config_name=semantic_config_name, + enable_caching=enable_caching, + cache_ttl_seconds=cache_ttl_seconds, + embedding_provider=embedding_provider, + embedding_model=embedding_model, + openai_api_key=openai_api_key, + openai_api_version=openai_api_version, + openai_endpoint=openai_endpoint, + ) + + self._endpoint = endpoint + self._index_name = index_name + self._credential = processed_credential + self._api_version = api_version + + self._client: Optional[SearchClient] = None + self._cache: Dict[str, Dict[str, Any]] = {} + + if self.search_config.api_version == "2023-11-01" and self.search_config.vector_fields: + warning_message = ( + f"When explicitly setting api_version='{self.search_config.api_version}' for vector search: " + f"If client-side embedding is NOT configured (e.g., `embedding_model` is not set), " + f"this tool defaults to service-side vectorization (VectorizableTextQuery), which may fail or have limitations with this API version. " + f"If client-side embedding IS configured, the tool will use VectorizedQuery, which is generally compatible. " + f"For robust vector search, consider omitting api_version (recommended to use SDK default) or use a newer API version." + ) + logger.warning(warning_message) + + async def close(self) -> None: + """Explicitly close the Azure SearchClient if needed (for cleanup).""" + if self._client is not None: + try: + await self._client.close() + except Exception: + pass + finally: + self._client = None + + def _process_credential( + self, credential: Union[AzureKeyCredential, AsyncTokenCredential, Dict[str, str]] + ) -> Union[AzureKeyCredential, AsyncTokenCredential]: + """Process credential to ensure it's the correct type for async SearchClient. + + Converts dictionary credentials with 'api_key' to AzureKeyCredential objects. + + Args: + credential: The credential in either object or dictionary form + + Returns: + A properly formatted credential object + + Raises: + ValueError: If the credential dictionary doesn't contain an 'api_key' + TypeError: If the credential is not of a supported type + """ + if isinstance(credential, dict): + if "api_key" in credential: + return AzureKeyCredential(credential["api_key"]) + raise ValueError("If credential is a dict, it must contain an 'api_key' key") + + if isinstance(credential, (AzureKeyCredential, AsyncTokenCredential)): + return credential + + raise TypeError("Credential must be AzureKeyCredential, AsyncTokenCredential, or a valid dict") + + async def _get_client(self) -> SearchClient: + """Get the search client for the configured index. + + Returns: + SearchClient: Initialized search client + + Raises: + ValueError: If index doesn't exist or authentication fails + """ + if self._client is not None: + return self._client + + try: + self._client = SearchClient( + endpoint=self.search_config.endpoint, + index_name=self.search_config.index_name, + credential=self.search_config.credential, + api_version=self.search_config.api_version, + ) + return self._client + except ResourceNotFoundError as e: + raise ValueError(f"Index '{self.search_config.index_name}' not found in Azure AI Search service.") from e + except HttpResponseError as e: + if e.status_code == 401: + raise ValueError("Authentication failed. Please check your credentials.") from e + elif e.status_code == 403: + raise ValueError("Permission denied to access this index.") from e + else: + raise ValueError(f"Error connecting to Azure AI Search: {str(e)}") from e + except Exception as e: + raise ValueError(f"Unexpected error initializing search client: {str(e)}") from e + + async def run( + self, args: Union[str, Dict[str, Any], SearchQuery], cancellation_token: Optional[CancellationToken] = None + ) -> SearchResults: + """Execute a search against the Azure AI Search index. + + Args: + args: Search query text or SearchQuery object + cancellation_token: Optional token to cancel the operation + + Returns: + SearchResults: Container with search results and metadata + + Raises: + ValueError: If the search query is empty or invalid + ValueError: If there is an authentication error or other search issue + asyncio.CancelledError: If the operation is cancelled + """ + if isinstance(args, str): + if not args.strip(): + raise ValueError("Search query cannot be empty") + search_query = SearchQuery(query=args) + elif isinstance(args, dict) and "query" in args: + search_query = SearchQuery(query=args["query"]) + elif isinstance(args, SearchQuery): + search_query = args + else: + raise ValueError("Invalid search query format. Expected string, dict with 'query', or SearchQuery") + + if cancellation_token is not None and cancellation_token.is_cancelled(): + raise asyncio.CancelledError("Operation cancelled") + + cache_key = "" + if self.search_config.enable_caching: + cache_key_parts = [ + search_query.query, + str(self.search_config.top), + self.search_config.query_type, + ",".join(sorted(self.search_config.search_fields or [])), + ",".join(sorted(self.search_config.select_fields or [])), + ",".join(sorted(self.search_config.vector_fields or [])), + str(self.search_config.filter or ""), + str(self.search_config.semantic_config_name or ""), + ] + cache_key = ":".join(filter(None, cache_key_parts)) + if cache_key in self._cache: + cache_entry = self._cache[cache_key] + cache_age = time.time() - cache_entry["timestamp"] + if cache_age < self.search_config.cache_ttl_seconds: + logger.debug(f"Using cached results for query: {search_query.query}") + return SearchResults( + results=[ + SearchResult(score=r.score, content=r.content, metadata=r.metadata) + for r in cache_entry["results"] + ] + ) + + try: + search_kwargs: Dict[str, Any] = {} + + if self.search_config.query_type != "vector": + search_kwargs["search_text"] = search_query.query + search_kwargs["query_type"] = self.search_config.query_type + + if self.search_config.search_fields: + search_kwargs["search_fields"] = self.search_config.search_fields # type: ignore[assignment] + + if self.search_config.query_type == "semantic" and self.search_config.semantic_config_name: + search_kwargs["semantic_configuration_name"] = self.search_config.semantic_config_name + + if self.search_config.select_fields: + search_kwargs["select"] = self.search_config.select_fields # type: ignore[assignment] + if self.search_config.filter: + search_kwargs["filter"] = str(self.search_config.filter) + if self.search_config.top is not None: + search_kwargs["top"] = self.search_config.top # type: ignore[assignment] + + if self.search_config.vector_fields and len(self.search_config.vector_fields) > 0: + if not search_query.query: + raise ValueError("Query text cannot be empty for vector search operations") + + use_client_side_embeddings = bool( + self.search_config.embedding_model and self.search_config.embedding_provider + ) + + vector_queries: List[Union[VectorizedQuery, VectorizableTextQuery]] = [] + if use_client_side_embeddings: + from azure.search.documents.models import VectorizedQuery + + embedding_vector: List[float] = await self._get_embedding(search_query.query) + for field_spec in self.search_config.vector_fields: + fields = field_spec if isinstance(field_spec, str) else ",".join(field_spec) + vector_queries.append( + VectorizedQuery( + vector=embedding_vector, + k_nearest_neighbors=self.search_config.top or 5, + fields=fields, + kind="vector", + ) + ) + else: + from azure.search.documents.models import VectorizableTextQuery + + for field in self.search_config.vector_fields: + fields = field if isinstance(field, str) else ",".join(field) + vector_queries.append( + VectorizableTextQuery( # type: ignore + text=search_query.query, + k_nearest_neighbors=self.search_config.top or 5, + fields=fields, + kind="vectorizable", + ) + ) + + search_kwargs["vector_queries"] = vector_queries # type: ignore[assignment] + + if cancellation_token is not None: + dummy_task = asyncio.create_task(asyncio.sleep(60)) + cancellation_token.link_future(dummy_task) + + def is_cancelled() -> bool: + return cancellation_token.is_cancelled() + else: + + def is_cancelled() -> bool: + return False + + client = await self._get_client() + search_results: SearchResultsIterable = await client.search(**search_kwargs) # type: ignore[arg-type] + + results: List[SearchResult] = [] + async for doc in search_results: + if is_cancelled(): + raise asyncio.CancelledError("Operation was cancelled") + + try: + metadata: Dict[str, Any] = {} + content: Dict[str, Any] = {} + + for key, value in doc.items(): + if isinstance(key, str) and key.startswith(("@", "_")): + metadata[key] = value + else: + content[str(key)] = value + + score = float(metadata.get("@search.score", 0.0)) + results.append(SearchResult(score=score, content=content, metadata=metadata)) + except Exception as e: + logger.warning(f"Error processing search document: {e}") + continue + + if self.search_config.enable_caching: + self._cache[cache_key] = {"results": results, "timestamp": time.time()} + + return SearchResults(results=results) + + except asyncio.CancelledError: + raise + except Exception as e: + error_msg = str(e) + if isinstance(e, HttpResponseError): + if hasattr(e, "message") and e.message: + error_msg = e.message + + if "not found" in error_msg.lower(): + raise ValueError(f"Index '{self.search_config.index_name}' not found.") from e + elif "unauthorized" in error_msg.lower() or "401" in error_msg: + raise ValueError(f"Authentication failed: {error_msg}") from e + else: + raise ValueError(f"Error from Azure AI Search: {error_msg}") from e + + def _to_config(self) -> AzureAISearchConfig: + """Convert the current instance to a configuration object.""" + return self.search_config + + @property + def schema(self) -> ToolSchema: + """Return the schema for the tool.""" + return { + "name": self.name, + "description": self.description, + "parameters": { + "type": "object", + "properties": {"query": {"type": "string", "description": "Search query text"}}, + "required": ["query"], + "additionalProperties": False, + }, + "strict": True, + } + + def return_value_as_string(self, value: SearchResults) -> str: + """Convert the search results to a string representation.""" + if not value.results: + return "No results found." + + result_strings: List[str] = [] + for i, result in enumerate(value.results, 1): + content_items = [f"{k}: {str(v) if v is not None else 'None'}" for k, v in result.content.items()] + content_str = ", ".join(content_items) + result_strings.append(f"Result {i} (Score: {result.score:.2f}): {content_str}") + + return "\n".join(result_strings) + + @classmethod + def _validate_config( + cls, config_dict: Dict[str, Any], search_type: Literal["full_text", "vector", "hybrid"] + ) -> None: + """Validate configuration for specific search types.""" + credential = config_dict.get("credential") + if isinstance(credential, str): + raise TypeError("Credential must be AzureKeyCredential, AsyncTokenCredential, or a valid dict") + if isinstance(credential, dict) and "api_key" not in credential: + raise ValueError("If credential is a dict, it must contain an 'api_key' key") + + try: + _ = AzureAISearchConfig(**config_dict) + except Exception as e: + raise ValueError(f"Invalid configuration: {str(e)}") from e + + if search_type == "vector": + vector_fields = config_dict.get("vector_fields") + if not vector_fields or len(vector_fields) == 0: + raise ValueError("vector_fields must contain at least one field name for vector search") + + elif search_type == "hybrid": + vector_fields = config_dict.get("vector_fields") + search_fields = config_dict.get("search_fields") + + if not vector_fields or len(vector_fields) == 0: + raise ValueError("vector_fields must contain at least one field name for hybrid search") + + if not search_fields or len(search_fields) == 0: + raise ValueError("search_fields must contain at least one field name for hybrid search") + + @classmethod + @abstractmethod + def _from_config(cls, config: AzureAISearchConfig) -> "BaseAzureAISearchTool": + """Create a tool instance from a configuration object. + + This is an abstract method that must be implemented by subclasses. + """ + if cls is BaseAzureAISearchTool: + raise NotImplementedError( + "BaseAzureAISearchTool is an abstract base class and cannot be instantiated directly. " + "Use a concrete implementation like AzureAISearchTool." + ) + raise NotImplementedError("Subclasses must implement _from_config") + + @abstractmethod + async def _get_embedding(self, query: str) -> List[float]: + """Generate embedding vector for the query text.""" + raise NotImplementedError("Subclasses must implement _get_embedding") + + +_allow_private_constructor = ContextVar("_allow_private_constructor", default=False) + + +class AzureAISearchTool(EmbeddingProviderMixin, BaseAzureAISearchTool): + """Azure AI Search tool for querying Azure search indexes. + + This tool provides a simplified interface for querying Azure AI Search indexes using + various search methods. It's recommended to use the factory methods to create + instances tailored for specific search types: + + 1. **Full-Text Search**: For traditional keyword-based searches, Lucene queries, or + semantically re-ranked results. + - Use `AzureAISearchTool.create_full_text_search()` + - Supports `query_type`: "simple" (keyword), "full" (Lucene), "semantic". + + 2. **Vector Search**: For pure similarity searches based on vector embeddings. + - Use `AzureAISearchTool.create_vector_search()` + + 3. **Hybrid Search**: For combining vector search with full-text or semantic search + to get the benefits of both. + - Use `AzureAISearchTool.create_hybrid_search()` + - The text component can be "simple", "full", or "semantic" via the `query_type` parameter. + + Each factory method configures the tool with appropriate defaults and validations + for the chosen search strategy. + + .. warning:: + If you set `query_type="semantic"`, you must also provide a valid `semantic_config_name`. + This configuration must be set up in your Azure AI Search index beforehand. + """ + + component_provider_override = "autogen_ext.tools.azure.AzureAISearchTool" + + @classmethod + def _from_config(cls, config: AzureAISearchConfig) -> "AzureAISearchTool": + """Create a tool instance from a configuration object. + + Args: + config: The configuration object with tool settings + + Returns: + AzureAISearchTool: An initialized tool instance + """ + token = _allow_private_constructor.set(True) + try: + instance = cls( + name=config.name, + description=config.description or "", + endpoint=config.endpoint, + index_name=config.index_name, + credential=config.credential, + api_version=config.api_version, + query_type=config.query_type, + search_fields=config.search_fields, + select_fields=config.select_fields, + vector_fields=config.vector_fields, + top=config.top, + filter=config.filter, + semantic_config_name=config.semantic_config_name, + enable_caching=config.enable_caching, + cache_ttl_seconds=config.cache_ttl_seconds, + embedding_provider=config.embedding_provider, + embedding_model=config.embedding_model, + openai_api_key=config.openai_api_key, + openai_api_version=config.openai_api_version, + openai_endpoint=config.openai_endpoint, + ) + return instance + finally: + _allow_private_constructor.reset(token) + + @classmethod + def _create_from_params( + cls, config_dict: Dict[str, Any], search_type: Literal["full_text", "vector", "hybrid"] + ) -> "AzureAISearchTool": + """Private helper to create an instance from parameters after validation. + + Args: + config_dict: Dictionary with configuration parameters + search_type: Type of search for validation + + Returns: + Configured AzureAISearchTool instance + """ + cls._validate_config(config_dict, search_type) + + token = _allow_private_constructor.set(True) + try: + return cls(**config_dict) + finally: + _allow_private_constructor.reset(token) + + @classmethod + def create_full_text_search( + cls, + name: str, + endpoint: str, + index_name: str, + credential: Union[AzureKeyCredential, AsyncTokenCredential, Dict[str, str]], + description: Optional[str] = None, + api_version: Optional[str] = None, + query_type: Literal["simple", "full", "semantic"] = "simple", + search_fields: Optional[List[str]] = None, + select_fields: Optional[List[str]] = None, + top: Optional[int] = 5, + filter: Optional[str] = None, + semantic_config_name: Optional[str] = None, + enable_caching: bool = False, + cache_ttl_seconds: int = 300, + ) -> "AzureAISearchTool": + """Create a tool for traditional text-based searches. + + This factory method creates an AzureAISearchTool optimized for full-text search, + supporting keyword matching, Lucene syntax, and semantic search capabilities. + + Args: + name: The name of this tool instance + endpoint: The full URL of your Azure AI Search service + index_name: Name of the search index to query + credential: Azure credential for authentication (API key or token) + description: Optional description explaining the tool's purpose + api_version: Azure AI Search API version to use + query_type: Type of text search to perform: + + • **simple** : Basic keyword search that matches exact terms and their variations + • **full**: Advanced search using Lucene query syntax for complex queries + • **semantic**: AI-powered search that understands meaning and context, providing enhanced relevance ranking + search_fields: Fields to search within documents + select_fields: Fields to return in search results + top: Maximum number of results to return (default: 5) + filter: OData filter expression to refine search results + semantic_config_name: Semantic configuration name (required for semantic query_type) + enable_caching: Whether to cache search results + cache_ttl_seconds: How long to cache results in seconds + + Returns: + An initialized AzureAISearchTool for full-text search + + Example: + .. code-block:: python + + from azure.core.credentials import AzureKeyCredential + from autogen_ext.tools.azure import AzureAISearchTool + + # Basic keyword search + tool = AzureAISearchTool.create_full_text_search( + name="doc-search", + endpoint="https://your-search.search.windows.net", # Your Azure AI Search endpoint + index_name="", # Name of your search index + credential=AzureKeyCredential(""), # Your Azure AI Search admin key + query_type="simple", # Enable keyword search + search_fields=["content", "title"], # Required: fields to search within + select_fields=["content", "title", "url"], # Optional: fields to return + top=5, + ) + + # full text (Lucene query) search + full_text_tool = AzureAISearchTool.create_full_text_search( + name="doc-search", + endpoint="https://your-search.search.windows.net", # Your Azure AI Search endpoint + index_name="", # Name of your search index + credential=AzureKeyCredential(""), # Your Azure AI Search admin key + query_type="full", # Enable Lucene query syntax + search_fields=["content", "title"], # Required: fields to search within + select_fields=["content", "title", "url"], # Optional: fields to return + top=5, + ) + + # Semantic search with re-ranking + # Note: Make sure your index has semantic configuration enabled + semantic_tool = AzureAISearchTool.create_full_text_search( + name="semantic-search", + endpoint="https://your-search.search.windows.net", + index_name="", + credential=AzureKeyCredential(""), + query_type="semantic", # Enable semantic ranking + semantic_config_name="", # Required for semantic search + search_fields=["content", "title"], # Required: fields to search within + select_fields=["content", "title", "url"], # Optional: fields to return + top=5, + ) + + # The search tool can be used with an Agent + # assistant = Agent("assistant", tools=[semantic_tool]) + """ + if query_type == "semantic" and not semantic_config_name: + raise ValueError("semantic_config_name is required when query_type is 'semantic'") + + config_dict = { + "name": name, + "endpoint": endpoint, + "index_name": index_name, + "credential": credential, + "description": description, + "api_version": api_version or DEFAULT_API_VERSION, + "query_type": query_type, + "search_fields": search_fields, + "select_fields": select_fields, + "top": top, + "filter": filter, + "semantic_config_name": semantic_config_name, + "enable_caching": enable_caching, + "cache_ttl_seconds": cache_ttl_seconds, + } + + return cls._create_from_params(config_dict, "full_text") + + @classmethod + def create_vector_search( + cls, + name: str, + endpoint: str, + index_name: str, + credential: Union[AzureKeyCredential, AsyncTokenCredential, Dict[str, str]], + vector_fields: List[str], + description: Optional[str] = None, + api_version: Optional[str] = None, + select_fields: Optional[List[str]] = None, + top: int = 5, + filter: Optional[str] = None, + enable_caching: bool = False, + cache_ttl_seconds: int = 300, + embedding_provider: Optional[str] = None, + embedding_model: Optional[str] = None, + openai_api_key: Optional[str] = None, + openai_api_version: Optional[str] = None, + openai_endpoint: Optional[str] = None, + ) -> "AzureAISearchTool": + """Create a tool for pure vector/similarity search. + + This factory method creates an AzureAISearchTool optimized for vector search, + allowing for semantic similarity-based matching using vector embeddings. + + Args: + name: The name of this tool instance + endpoint: The full URL of your Azure AI Search service + index_name: Name of the search index to query + credential: Azure credential for authentication (API key or token) + vector_fields: Fields to use for vector search (required) + description: Optional description explaining the tool's purpose + api_version: Azure AI Search API version to use + select_fields: Fields to return in search results + top: Maximum number of results to return / k in k-NN (default: 5) + filter: OData filter expression to refine search results + enable_caching: Whether to cache search results + cache_ttl_seconds: How long to cache results in seconds + embedding_provider: Provider for client-side embeddings (e.g., 'azure_openai', 'openai') + embedding_model: Model for client-side embeddings (e.g., 'text-embedding-ada-002') + openai_api_key: API key for OpenAI/Azure OpenAI embeddings + openai_api_version: API version for Azure OpenAI embeddings + openai_endpoint: Endpoint URL for Azure OpenAI embeddings + + Returns: + An initialized AzureAISearchTool for vector search + + Raises: + ValueError: If vector_fields is empty + ValueError: If embedding_provider is 'azure_openai' without openai_endpoint + ValueError: If required parameters are missing or invalid + + Example Usage: + .. code-block:: python + + from azure.core.credentials import AzureKeyCredential + from autogen_ext.tools.azure import AzureAISearchTool + + # Vector search with service-side vectorization + tool = AzureAISearchTool.create_vector_search( + name="vector-search", + endpoint="https://your-search.search.windows.net", # Your Azure AI Search endpoint + index_name="", # Name of your search index + credential=AzureKeyCredential(""), # Your Azure AI Search admin key + vector_fields=["content_vector"], # Your vector field name + select_fields=["content", "title", "url"], # Fields to return in results + top=5, + ) + + # Vector search with Azure OpenAI embeddings + azure_openai_tool = AzureAISearchTool.create_vector_search( + name="azure-openai-vector-search", + endpoint="https://your-search.search.windows.net", + index_name="", + credential=AzureKeyCredential(""), + vector_fields=["content_vector"], + embedding_provider="azure_openai", # Use Azure OpenAI for embeddings + embedding_model="text-embedding-ada-002", # Embedding model to use + openai_endpoint="https://your-openai.openai.azure.com", # Your Azure OpenAI endpoint + openai_api_key="", # Your Azure OpenAI key + openai_api_version="2024-02-15-preview", # Azure OpenAI API version + select_fields=["content", "title", "url"], # Fields to return in results + top=5, + ) + + # Vector search with OpenAI embeddings + openai_tool = AzureAISearchTool.create_vector_search( + name="openai-vector-search", + endpoint="https://your-search.search.windows.net", + index_name="", + credential=AzureKeyCredential(""), + vector_fields=["content_vector"], + embedding_provider="openai", # Use OpenAI for embeddings + embedding_model="text-embedding-ada-002", # Embedding model to use + openai_api_key="", # Your OpenAI API key + select_fields=["content", "title", "url"], # Fields to return in results + top=5, + ) + + # Use the tool with an Agent + # assistant = Agent("assistant", tools=[azure_openai_tool]) + """ + if embedding_provider == "azure_openai" and not openai_endpoint: + raise ValueError("openai_endpoint is required when embedding_provider is 'azure_openai'") + + config_dict = { + "name": name, + "endpoint": endpoint, + "index_name": index_name, + "credential": credential, + "description": description, + "api_version": api_version or DEFAULT_API_VERSION, + "query_type": "vector", + "select_fields": select_fields, + "vector_fields": vector_fields, + "top": top, + "filter": filter, + "enable_caching": enable_caching, + "cache_ttl_seconds": cache_ttl_seconds, + "embedding_provider": embedding_provider, + "embedding_model": embedding_model, + "openai_api_key": openai_api_key, + "openai_api_version": openai_api_version, + "openai_endpoint": openai_endpoint, + } + + return cls._create_from_params(config_dict, "vector") + + @classmethod + def create_hybrid_search( + cls, + name: str, + endpoint: str, + index_name: str, + credential: Union[AzureKeyCredential, AsyncTokenCredential, Dict[str, str]], + vector_fields: List[str], + search_fields: List[str], + description: Optional[str] = None, + api_version: Optional[str] = None, + query_type: Literal["simple", "full", "semantic"] = "simple", + select_fields: Optional[List[str]] = None, + top: int = 5, + filter: Optional[str] = None, + semantic_config_name: Optional[str] = None, + enable_caching: bool = False, + cache_ttl_seconds: int = 300, + embedding_provider: Optional[str] = None, + embedding_model: Optional[str] = None, + openai_api_key: Optional[str] = None, + openai_api_version: Optional[str] = None, + openai_endpoint: Optional[str] = None, + ) -> "AzureAISearchTool": + """Create a tool that combines vector and text search capabilities. + + This factory method creates an AzureAISearchTool configured for hybrid search, + which combines the benefits of vector similarity and traditional text search. + + Args: + name: The name of this tool instance + endpoint: The full URL of your Azure AI Search service + index_name: Name of the search index to query + credential: Azure credential for authentication (API key or token) + vector_fields: Fields to use for vector search (required) + search_fields: Fields to use for text search (required) + description: Optional description explaining the tool's purpose + api_version: Azure AI Search API version to use + query_type: Type of text search to perform: + + • **simple**: Basic keyword search that matches exact terms and their variations + • **full**: Advanced search using Lucene query syntax for complex queries + • **semantic**: AI-powered search that understands meaning and context, providing enhanced relevance ranking + select_fields: Fields to return in search results + top: Maximum number of results to return (default: 5) + filter: OData filter expression to refine search results + semantic_config_name: Semantic configuration name (required if query_type="semantic") + enable_caching: Whether to cache search results + cache_ttl_seconds: How long to cache results in seconds + embedding_provider: Provider for client-side embeddings (e.g., 'azure_openai', 'openai') + embedding_model: Model for client-side embeddings (e.g., 'text-embedding-ada-002') + openai_api_key: API key for OpenAI/Azure OpenAI embeddings + openai_api_version: API version for Azure OpenAI embeddings + openai_endpoint: Endpoint URL for Azure OpenAI embeddings + + Returns: + An initialized AzureAISearchTool for hybrid search + + Raises: + ValueError: If vector_fields or search_fields is empty + ValueError: If query_type is "semantic" without semantic_config_name + ValueError: If embedding_provider is 'azure_openai' without openai_endpoint + ValueError: If required parameters are missing or invalid + + Example: + .. code-block:: python + + from azure.core.credentials import AzureKeyCredential + from autogen_ext.tools.azure import AzureAISearchTool + + # Basic hybrid search with service-side vectorization + tool = AzureAISearchTool.create_hybrid_search( + name="hybrid-search", + endpoint="https://your-search.search.windows.net", # Your Azure AI Search endpoint + index_name="", # Name of your search index + credential=AzureKeyCredential(""), # Your Azure AI Search admin key + vector_fields=["content_vector"], # Your vector field name + search_fields=["content", "title"], # Your searchable fields + top=5, + ) + + # Hybrid search with semantic ranking and Azure OpenAI embeddings + semantic_tool = AzureAISearchTool.create_hybrid_search( + name="semantic-hybrid-search", + endpoint="https://your-search.search.windows.net", + index_name="", + credential=AzureKeyCredential(""), + vector_fields=["content_vector"], + search_fields=["content", "title"], + query_type="semantic", # Enable semantic ranking + semantic_config_name="", # Your semantic config name + embedding_provider="azure_openai", # Use Azure OpenAI for embeddings + embedding_model="text-embedding-ada-002", # Embedding model to use + openai_endpoint="https://your-openai.openai.azure.com", # Your Azure OpenAI endpoint + openai_api_key="", # Your Azure OpenAI key + openai_api_version="2024-02-15-preview", # Azure OpenAI API version + select_fields=["content", "title", "url"], # Fields to return in results + filter="language eq 'en'", # Optional OData filter + top=5, + ) + + # The search tool can be used with an Agent + # assistant = Agent("assistant", tools=[semantic_tool]) + """ + if query_type == "semantic" and not semantic_config_name: + raise ValueError("semantic_config_name is required when query_type is 'semantic'") + + if embedding_provider == "azure_openai" and not openai_endpoint: + raise ValueError("openai_endpoint is required when embedding_provider is 'azure_openai'") + + config_dict = { + "name": name, + "endpoint": endpoint, + "index_name": index_name, + "credential": credential, + "description": description, + "api_version": api_version or DEFAULT_API_VERSION, + "query_type": query_type, + "search_fields": search_fields, + "select_fields": select_fields, + "vector_fields": vector_fields, + "top": top, + "filter": filter, + "semantic_config_name": semantic_config_name, + "enable_caching": enable_caching, + "cache_ttl_seconds": cache_ttl_seconds, + "embedding_provider": embedding_provider, + "embedding_model": embedding_model, + "openai_api_key": openai_api_key, + "openai_api_version": openai_api_version, + "openai_endpoint": openai_endpoint, + } + + return cls._create_from_params(config_dict, "hybrid") diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/azure/_config.py b/python/packages/autogen-ext/src/autogen_ext/tools/azure/_config.py new file mode 100644 index 000000000000..8af10c1d8852 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/tools/azure/_config.py @@ -0,0 +1,186 @@ +"""Configuration for Azure AI Search tool. + +This module provides configuration classes for the Azure AI Search tool, including +settings for authentication, search behavior, retry policies, and caching. +""" + +import logging +from typing import ( + List, + Literal, + Optional, + TypeVar, + Union, +) + +from pydantic import BaseModel, Field, field_validator, model_validator + +from azure.core.credentials import AzureKeyCredential +from azure.core.credentials_async import AsyncTokenCredential + +T = TypeVar("T", bound="AzureAISearchConfig") + +logger = logging.getLogger(__name__) + +QueryTypeLiteral = Literal["simple", "full", "semantic", "vector"] +DEFAULT_API_VERSION = "2023-10-01-preview" + + +class AzureAISearchConfig(BaseModel): + """Configuration for Azure AI Search with validation. + + This class defines the configuration parameters for Azure AI Search tools, including + authentication, search behavior, caching, and embedding settings. + + .. note:: + This class requires the ``azure`` extra for the ``autogen-ext`` package. + + .. code-block:: bash + + pip install -U "autogen-ext[azure]" + + .. note:: + **Prerequisites:** + + 1. An Azure AI Search service must be created in your Azure subscription. + 2. The search index must be properly configured for your use case: + + - For vector search: Index must have vector fields + - For semantic search: Index must have semantic configuration + - For hybrid search: Both vector fields and text fields must be configured + 3. Required packages: + + - Base functionality: ``azure-search-documents>=11.4.0`` + - For Azure OpenAI embeddings: ``openai azure-identity`` + - For OpenAI embeddings: ``openai`` + + Example Usage: + .. code-block:: python + + from azure.core.credentials import AzureKeyCredential + from autogen_ext.tools.azure import AzureAISearchConfig + + # Basic configuration for full-text search + config = AzureAISearchConfig( + name="doc-search", + endpoint="https://your-search.search.windows.net", # Your Azure AI Search endpoint + index_name="", # Name of your search index + credential=AzureKeyCredential(""), # Your Azure AI Search admin key + query_type="simple", + search_fields=["content", "title"], # Update with your searchable fields + top=5, + ) + + # Configuration for vector search with Azure OpenAI embeddings + vector_config = AzureAISearchConfig( + name="vector-search", + endpoint="https://your-search.search.windows.net", + index_name="", + credential=AzureKeyCredential(""), + query_type="vector", + vector_fields=["embedding"], # Update with your vector field name + embedding_provider="azure_openai", + embedding_model="text-embedding-ada-002", + openai_endpoint="https://your-openai.openai.azure.com", # Your Azure OpenAI endpoint + openai_api_key="", # Your Azure OpenAI key + top=5, + ) + + # Configuration for hybrid search with semantic ranking + hybrid_config = AzureAISearchConfig( + name="hybrid-search", + endpoint="https://your-search.search.windows.net", + index_name="", + credential=AzureKeyCredential(""), + query_type="semantic", + semantic_config_name="", # Name of your semantic configuration + search_fields=["content", "title"], # Update with your search fields + vector_fields=["embedding"], # Update with your vector field name + embedding_provider="openai", + embedding_model="text-embedding-ada-002", + openai_api_key="", # Your OpenAI API key + top=5, + ) + """ + + name: str = Field(description="The name of this tool instance") + description: Optional[str] = Field(default=None, description="Description explaining the tool's purpose") + endpoint: str = Field(description="The full URL of your Azure AI Search service") + index_name: str = Field(description="Name of the search index to query") + credential: Union[AzureKeyCredential, AsyncTokenCredential] = Field( + description="Azure credential for authentication (API key or token)" + ) + api_version: str = Field( + default=DEFAULT_API_VERSION, + description=f"Azure AI Search API version to use. Defaults to {DEFAULT_API_VERSION}.", + ) + query_type: QueryTypeLiteral = Field( + default="simple", description="Type of search to perform: simple, full, semantic, or vector" + ) + search_fields: Optional[List[str]] = Field(default=None, description="Fields to search within documents") + select_fields: Optional[List[str]] = Field(default=None, description="Fields to return in search results") + vector_fields: Optional[List[str]] = Field(default=None, description="Fields to use for vector search") + top: Optional[int] = Field( + default=None, description="Maximum number of results to return. For vector searches, acts as k in k-NN." + ) + filter: Optional[str] = Field(default=None, description="OData filter expression to refine search results") + semantic_config_name: Optional[str] = Field( + default=None, description="Semantic configuration name for enhanced results" + ) + + enable_caching: bool = Field(default=False, description="Whether to cache search results") + cache_ttl_seconds: int = Field(default=300, description="How long to cache results in seconds") + + embedding_provider: Optional[str] = Field( + default=None, description="Name of embedding provider for client-side embeddings" + ) + embedding_model: Optional[str] = Field(default=None, description="Model name for client-side embeddings") + openai_api_key: Optional[str] = Field(default=None, description="API key for OpenAI/Azure OpenAI embeddings") + openai_api_version: Optional[str] = Field(default=None, description="API version for Azure OpenAI embeddings") + openai_endpoint: Optional[str] = Field(default=None, description="Endpoint URL for Azure OpenAI embeddings") + + model_config = {"arbitrary_types_allowed": True} + + @field_validator("endpoint") + def validate_endpoint(cls, v: str) -> str: + """Validate that the endpoint is a valid URL.""" + if not v.startswith(("http://", "https://")): + raise ValueError("endpoint must be a valid URL starting with http:// or https://") + return v + + @field_validator("query_type") + def normalize_query_type(cls, v: QueryTypeLiteral) -> QueryTypeLiteral: + """Normalize query type to standard values.""" + if not v: + return "simple" + + if isinstance(v, str) and v.lower() == "fulltext": + return "full" + + return v + + @field_validator("top") + def validate_top(cls, v: Optional[int]) -> Optional[int]: + """Ensure top is a positive integer if provided.""" + if v is not None and v <= 0: + raise ValueError("top must be a positive integer") + return v + + @model_validator(mode="after") + def validate_interdependent_fields(self) -> "AzureAISearchConfig": + """Validate interdependent fields after all fields have been parsed.""" + if self.query_type == "semantic" and not self.semantic_config_name: + raise ValueError("semantic_config_name must be provided when query_type is 'semantic'") + + if self.query_type == "vector" and not self.vector_fields: + raise ValueError("vector_fields must be provided for vector search") + + if ( + self.embedding_provider + and self.embedding_provider.lower() == "azure_openai" + and self.embedding_model + and not self.openai_endpoint + ): + raise ValueError("openai_endpoint must be provided for azure_openai embedding provider") + + return self diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_config.py b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_config.py index 10c3d4a985e6..b7df432ff856 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_config.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_config.py @@ -3,19 +3,19 @@ class DataConfig(BaseModel): input_dir: str - entity_table: str = "create_final_nodes" - entity_embedding_table: str = "create_final_entities" + entity_table: str = "entities" + entity_embedding_table: str = "entities" + community_table: str = "communities" community_level: int = 2 class GlobalDataConfig(DataConfig): - community_table: str = "create_final_communities" - community_report_table: str = "create_final_community_reports" + community_report_table: str = "community_reports" class LocalDataConfig(DataConfig): - relationship_table: str = "create_final_relationships" - text_unit_table: str = "create_final_text_units" + relationship_table: str = "relationships" + text_unit_table: str = "text_units" class ContextConfig(BaseModel): diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_global_search.py b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_global_search.py index 1a8915ee9bc9..937cce05f5e0 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_global_search.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_global_search.py @@ -1,21 +1,22 @@ -# mypy: disable-error-code="no-any-unimported,misc" from pathlib import Path import pandas as pd import tiktoken from autogen_core import CancellationToken from autogen_core.tools import BaseTool -from graphrag.config.config_file_loader import load_config_from_file +from pydantic import BaseModel, Field + +import graphrag.config.defaults as defs +from graphrag.config.load_config import load_config +from graphrag.language_model.manager import ModelManager +from graphrag.language_model.protocol import ChatModel from graphrag.query.indexer_adapters import ( read_indexer_communities, read_indexer_entities, read_indexer_reports, ) -from graphrag.query.llm.base import BaseLLM -from graphrag.query.llm.get_client import get_llm from graphrag.query.structured_search.global_search.community_context import GlobalCommunityContext from graphrag.query.structured_search.global_search.search import GlobalSearch -from pydantic import BaseModel, Field from ._config import GlobalContextConfig as ContextConfig from ._config import GlobalDataConfig as DataConfig @@ -63,6 +64,7 @@ class GlobalSearchTool(BaseTool[GlobalSearchToolArgs, GlobalSearchToolReturn]): .. code-block:: python import asyncio + from pathlib import Path from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_agentchat.ui import Console from autogen_ext.tools.graphrag import GlobalSearchTool @@ -77,7 +79,7 @@ async def main(): ) # Set up global search tool - global_tool = GlobalSearchTool.from_settings(settings_path="./settings.yaml") + global_tool = GlobalSearchTool.from_settings(root_dir=Path("./"), config_filepath=Path("./settings.yaml")) # Create assistant agent with the global search tool assistant_agent = AssistantAgent( @@ -103,7 +105,7 @@ async def main(): def __init__( self, token_encoder: tiktoken.Encoding, - llm: BaseLLM, + model: ChatModel, data_config: DataConfig, context_config: ContextConfig = _default_context_config, mapreduce_config: MapReduceConfig = _default_mapreduce_config, @@ -114,8 +116,8 @@ def __init__( name="global_search_tool", description="Perform a global search with given parameters using graphrag.", ) - # Use the provided LLM - self._llm = llm + # Use the provided model + self._model = model # Load parquet files community_df: pd.DataFrame = pd.read_parquet(f"{data_config.input_dir}/{data_config.community_table}.parquet") # type: ignore @@ -123,13 +125,11 @@ def __init__( report_df: pd.DataFrame = pd.read_parquet( # type: ignore f"{data_config.input_dir}/{data_config.community_report_table}.parquet" ) - entity_embedding_df: pd.DataFrame = pd.read_parquet( # type: ignore - f"{data_config.input_dir}/{data_config.entity_embedding_table}.parquet" - ) - communities = read_indexer_communities(community_df, entity_df, report_df) - reports = read_indexer_reports(report_df, entity_df, data_config.community_level) - entities = read_indexer_entities(entity_df, entity_embedding_df, data_config.community_level) + # Fix: Use correct argument order and types for GraphRAG API + communities = read_indexer_communities(community_df, report_df) + reports = read_indexer_reports(report_df, community_df, data_config.community_level) + entities = read_indexer_entities(entity_df, community_df, data_config.community_level) context_builder = GlobalCommunityContext( community_reports=reports, @@ -163,7 +163,7 @@ def __init__( } self._search_engine = GlobalSearch( - llm=self._llm, + model=self._model, context_builder=context_builder, token_encoder=token_encoder, max_data_tokens=context_config.max_data_tokens, @@ -177,37 +177,56 @@ def __init__( ) async def run(self, args: GlobalSearchToolArgs, cancellation_token: CancellationToken) -> GlobalSearchToolReturn: - search_result = await self._search_engine.asearch(args.query) + search_result = await self._search_engine.search(args.query) assert isinstance(search_result.response, str), "Expected response to be a string" return GlobalSearchToolReturn(answer=search_result.response) @classmethod - def from_settings(cls, settings_path: str | Path) -> "GlobalSearchTool": + def from_settings(cls, root_dir: str | Path, config_filepath: str | Path | None = None) -> "GlobalSearchTool": """Create a GlobalSearchTool instance from GraphRAG settings file. Args: - settings_path: Path to the GraphRAG settings.yaml file + root_dir: Path to the GraphRAG root directory + config_filepath: Path to the GraphRAG settings file (optional) Returns: An initialized GlobalSearchTool instance """ # Load GraphRAG config - config = load_config_from_file(settings_path) - - # Initialize token encoder - token_encoder = tiktoken.get_encoding(config.encoding_model) - - # Initialize LLM using graphrag's get_client - llm = get_llm(config) + if isinstance(root_dir, str): + root_dir = Path(root_dir) + if isinstance(config_filepath, str): + config_filepath = Path(config_filepath) + config = load_config(root_dir=root_dir, config_filepath=config_filepath) + + # Get the language model configuration from the models section + chat_model_config = config.models.get(defs.DEFAULT_CHAT_MODEL_ID) + + if chat_model_config is None: + raise ValueError("default_chat_model not found in config.models") + + # Initialize token encoder based on the model being used + try: + token_encoder = tiktoken.encoding_for_model(chat_model_config.model) + except KeyError: + # Fallback to cl100k_base if model is not recognized by tiktoken + token_encoder = tiktoken.get_encoding("cl100k_base") + + # Create the LLM using ModelManager + model = ModelManager().get_or_create_chat_model( + name="global_search_model", + model_type=chat_model_config.type, + config=chat_model_config, + ) # Create data config from storage paths data_config = DataConfig( - input_dir=str(Path(config.storage.base_dir)), + input_dir=str(config.output.base_dir), ) return cls( token_encoder=token_encoder, - llm=llm, + model=model, data_config=data_config, context_config=_default_context_config, mapreduce_config=_default_mapreduce_config, diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_local_search.py b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_local_search.py index 625c7a4e1cc7..7c0420275f00 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_local_search.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_local_search.py @@ -1,23 +1,24 @@ # mypy: disable-error-code="no-any-unimported,misc" -import os from pathlib import Path import pandas as pd import tiktoken from autogen_core import CancellationToken from autogen_core.tools import BaseTool -from graphrag.config.config_file_loader import load_config_from_file +from pydantic import BaseModel, Field + +import graphrag.config.defaults as defs +from graphrag.config.load_config import load_config +from graphrag.language_model.manager import ModelManager +from graphrag.language_model.protocol import ChatModel, EmbeddingModel from graphrag.query.indexer_adapters import ( read_indexer_entities, read_indexer_relationships, read_indexer_text_units, ) -from graphrag.query.llm.base import BaseLLM, BaseTextEmbedding -from graphrag.query.llm.get_client import get_llm, get_text_embedder from graphrag.query.structured_search.local_search.mixed_context import LocalSearchMixedContext from graphrag.query.structured_search.local_search.search import LocalSearch from graphrag.vector_stores.lancedb import LanceDBVectorStore -from pydantic import BaseModel, Field from ._config import LocalContextConfig, SearchConfig from ._config import LocalDataConfig as DataConfig @@ -63,6 +64,7 @@ class LocalSearchTool(BaseTool[LocalSearchToolArgs, LocalSearchToolReturn]): .. code-block:: python import asyncio + from pathlib import Path from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_agentchat.ui import Console from autogen_ext.tools.graphrag import LocalSearchTool @@ -77,7 +79,7 @@ async def main(): ) # Set up local search tool - local_tool = LocalSearchTool.from_settings(settings_path="./settings.yaml") + local_tool = LocalSearchTool.from_settings(root_dir=Path("./"), config_filepath=Path("./settings.yaml")) # Create assistant agent with the local search tool assistant_agent = AssistantAgent( @@ -102,8 +104,8 @@ async def main(): Args: token_encoder (tiktoken.Encoding): The tokenizer used for text encoding - llm (BaseLLM): The language model to use for search - embedder (BaseTextEmbedding): The text embedding model to use + model: The chat model to use for search (GraphRAG ChatModel) + embedder: The text embedding model to use (GraphRAG EmbeddingModel) data_config (DataConfig): Configuration for data source locations and settings context_config (LocalContextConfig, optional): Configuration for context building. Defaults to default config. search_config (SearchConfig, optional): Configuration for search operations. Defaults to default config. @@ -112,8 +114,8 @@ async def main(): def __init__( self, token_encoder: tiktoken.Encoding, - llm: BaseLLM, - embedder: BaseTextEmbedding, + model: ChatModel, # ChatModel from GraphRAG + embedder: EmbeddingModel, # EmbeddingModel from GraphRAG data_config: DataConfig, context_config: LocalContextConfig = _default_context_config, search_config: SearchConfig = _default_search_config, @@ -124,30 +126,23 @@ def __init__( name="local_search_tool", description="Perform a local search with given parameters using graphrag.", ) - # Use the adapter - self._llm = llm + # Use the provided models + self._model = model self._embedder = embedder # Load parquet files entity_df: pd.DataFrame = pd.read_parquet(f"{data_config.input_dir}/{data_config.entity_table}.parquet") # type: ignore - entity_embedding_df: pd.DataFrame = pd.read_parquet( # type: ignore - f"{data_config.input_dir}/{data_config.entity_embedding_table}.parquet" - ) relationship_df: pd.DataFrame = pd.read_parquet( # type: ignore f"{data_config.input_dir}/{data_config.relationship_table}.parquet" ) text_unit_df: pd.DataFrame = pd.read_parquet(f"{data_config.input_dir}/{data_config.text_unit_table}.parquet") # type: ignore + community_df: pd.DataFrame = pd.read_parquet(f"{data_config.input_dir}/{data_config.community_table}.parquet") # type: ignore # Read data using indexer adapters - entities = read_indexer_entities(entity_df, entity_embedding_df, data_config.community_level) + entities = read_indexer_entities(entity_df, community_df, data_config.community_level) relationships = read_indexer_relationships(relationship_df) text_units = read_indexer_text_units(text_unit_df) # Set up vector store for entity embeddings - description_embedding_store = LanceDBVectorStore( - collection_name="default-entity-description", - ) - description_embedding_store.connect(db_uri=os.path.join(data_config.input_dir, "lancedb")) - description_embedding_store = LanceDBVectorStore( collection_name="default-entity-description", ) @@ -179,47 +174,70 @@ def __init__( } self._search_engine = LocalSearch( - llm=self._llm, + model=self._model, context_builder=context_builder, token_encoder=token_encoder, - llm_params=llm_params, - context_builder_params=context_builder_params, response_type=search_config.response_type, + context_builder_params=context_builder_params, + model_params=llm_params, ) async def run(self, args: LocalSearchToolArgs, cancellation_token: CancellationToken) -> LocalSearchToolReturn: - search_result = await self._search_engine.asearch(args.query) # type: ignore + search_result = await self._search_engine.search(args.query) # type: ignore[reportUnknownMemberType] assert isinstance(search_result.response, str), "Expected response to be a string" return LocalSearchToolReturn(answer=search_result.response) @classmethod - def from_settings(cls, settings_path: str | Path) -> "LocalSearchTool": + def from_settings(cls, root_dir: Path, config_filepath: Path | None = None) -> "LocalSearchTool": """Create a LocalSearchTool instance from GraphRAG settings file. Args: - settings_path: Path to the GraphRAG settings.yaml file + root_dir: Path to the GraphRAG root directory + config_filepath: Path to the GraphRAG settings file (optional) Returns: An initialized LocalSearchTool instance """ # Load GraphRAG config - config = load_config_from_file(settings_path) - - # Initialize token encoder - token_encoder = tiktoken.get_encoding(config.encoding_model) + config = load_config(root_dir=root_dir, config_filepath=config_filepath) + + # Get the language model configurations from the models section + chat_model_config = config.models.get(defs.DEFAULT_CHAT_MODEL_ID) + embedding_model_config = config.models.get(defs.DEFAULT_EMBEDDING_MODEL_ID) + + if chat_model_config is None: + raise ValueError("default_chat_model not found in config.models") + if embedding_model_config is None: + raise ValueError("default_embedding_model not found in config.models") + + # Initialize token encoder based on the model being used + try: + token_encoder = tiktoken.encoding_for_model(chat_model_config.model) + except KeyError: + # Fallback to cl100k_base if model is not recognized by tiktoken + token_encoder = tiktoken.get_encoding("cl100k_base") + + # Create the models using ModelManager + model = ModelManager().get_or_create_chat_model( + name="local_search_model", + model_type=chat_model_config.type, + config=chat_model_config, + ) - # Initialize LLM and embedder using graphrag's get_client functions - llm = get_llm(config) - embedder = get_text_embedder(config) + embedder = ModelManager().get_or_create_embedding_model( + name="local_search_embedder", + model_type=embedding_model_config.type, + config=embedding_model_config, + ) # Create data config from storage paths data_config = DataConfig( - input_dir=str(Path(config.storage.base_dir)), + input_dir=str(config.output.base_dir), ) return cls( token_encoder=token_encoder, - llm=llm, + model=model, embedder=embedder, data_config=data_config, context_config=_default_context_config, diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/http/_http_tool.py b/python/packages/autogen-ext/src/autogen_ext/tools/http/_http_tool.py index 451d5826bad7..c519143be4c5 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/http/_http_tool.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/http/_http_tool.py @@ -8,6 +8,8 @@ from pydantic import BaseModel, Field from typing_extensions import Self +DEFAULT_TIMEOUT_CONFIG = 5.0 + class HttpToolConfig(BaseModel): name: str @@ -53,6 +55,10 @@ class HttpToolConfig(BaseModel): """ The type of response to return from the tool. """ + timeout: float = DEFAULT_TIMEOUT_CONFIG + """ + The timeout for the tool request in seconds. + """ class HttpTool(BaseTool[BaseModel, Any], Component[HttpToolConfig]): @@ -73,6 +79,8 @@ class HttpTool(BaseTool[BaseModel, Any], Component[HttpToolConfig]): Path parameters must also be included in the schema and must be strings. return_type (Literal["text", "json"], optional): The type of response to return from the tool. Defaults to "text". + timeout (float, optional): The timeout for HTTP requests in seconds. + Defaults to 5.0. .. note:: This tool requires the :code:`http-tool` extra for the :code:`autogen-ext` package. @@ -126,7 +134,7 @@ async def main(): [TextMessage(content="Can you base64 decode the value 'YWJjZGU=', please?", source="user")], CancellationToken(), ) - print(response.chat_message.content) + print(response.chat_message) asyncio.run(main()) @@ -148,6 +156,7 @@ def __init__( scheme: Literal["http", "https"] = "http", method: Literal["GET", "POST", "PUT", "DELETE", "PATCH"] = "POST", return_type: Literal["text", "json"] = "text", + timeout: float = DEFAULT_TIMEOUT_CONFIG, ) -> None: self.server_params = HttpToolConfig( name=name, @@ -160,6 +169,7 @@ def __init__( headers=headers, json_schema=json_schema, return_type=return_type, + timeout=timeout, ) # Use regex to find all path parameters, we will need those later to template the path @@ -211,7 +221,8 @@ async def run(self, args: BaseModel, cancellation_token: CancellationToken) -> A port=self.server_params.port, path=path, ) - async with httpx.AsyncClient() as client: + timeout_config = httpx.Timeout(timeout=self.server_params.timeout) + async with httpx.AsyncClient(timeout=timeout_config) as client: match self.server_params.method: case "GET": response = await client.get(url, headers=self.server_params.headers, params=model_dump) diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/__init__.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/__init__.py index 83d76fcad502..fcf4148743a0 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/__init__.py @@ -1,13 +1,22 @@ -from ._config import McpServerParams, SseServerParams, StdioServerParams +from ._actor import McpSessionActor +from ._config import McpServerParams, SseServerParams, StdioServerParams, StreamableHttpServerParams from ._factory import mcp_server_tools +from ._session import create_mcp_server_session from ._sse import SseMcpToolAdapter from ._stdio import StdioMcpToolAdapter +from ._streamable_http import StreamableHttpMcpToolAdapter +from ._workbench import McpWorkbench __all__ = [ + "create_mcp_server_session", + "McpSessionActor", "StdioMcpToolAdapter", "StdioServerParams", "SseMcpToolAdapter", "SseServerParams", + "StreamableHttpMcpToolAdapter", + "StreamableHttpServerParams", "McpServerParams", "mcp_server_tools", + "McpWorkbench", ] diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_actor.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_actor.py new file mode 100644 index 000000000000..f909aee25cab --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_actor.py @@ -0,0 +1,310 @@ +import asyncio +import atexit +import base64 +import io +import logging +from typing import Any, Coroutine, Dict, Mapping, TypedDict + +from autogen_core import Component, ComponentBase, ComponentModel, Image +from autogen_core.models import ( + AssistantMessage, + ChatCompletionClient, + LLMMessage, + ModelInfo, + SystemMessage, + UserMessage, +) +from PIL import Image as PILImage +from pydantic import BaseModel +from typing_extensions import Self + +from mcp import types as mcp_types +from mcp.client.session import ClientSession +from mcp.shared.context import RequestContext + +from ._config import McpServerParams +from ._session import create_mcp_server_session + +logger = logging.getLogger(__name__) + +McpResult = ( + Coroutine[Any, Any, mcp_types.ListToolsResult] + | Coroutine[Any, Any, mcp_types.CallToolResult] + | Coroutine[Any, Any, mcp_types.ListPromptsResult] + | Coroutine[Any, Any, mcp_types.ListResourcesResult] + | Coroutine[Any, Any, mcp_types.ListResourceTemplatesResult] + | Coroutine[Any, Any, mcp_types.ReadResourceResult] + | Coroutine[Any, Any, mcp_types.GetPromptResult] +) +McpFuture = asyncio.Future[McpResult] + + +def _parse_sampling_content( + content: mcp_types.TextContent | mcp_types.ImageContent | mcp_types.AudioContent, model_info: ModelInfo +) -> str | Image: + """Convert MCP content types to Autogen content types.""" + if content.type == "text": + return content.text + elif content.type == "image": + if not model_info["vision"]: + raise ValueError("Sampling model does not support image content.") + # Decode base64 image data and create PIL Image + image_data = base64.b64decode(content.data) + pil_image = PILImage.open(io.BytesIO(image_data)) + return Image.from_pil(pil_image) + else: + raise ValueError(f"Unsupported content type: {content.type}") + + +def _parse_sampling_message(message: mcp_types.SamplingMessage, model_info: ModelInfo) -> LLMMessage: + """Convert MCP sampling messages to Autogen messages.""" + content = _parse_sampling_content(message.content, model_info=model_info) + if message.role == "user": + return UserMessage( + source="user", + content=[content], + ) + elif message.role == "assistant": + assert isinstance(content, str), "Assistant messages only support string content." + return AssistantMessage( + source="assistant", + content=content, + ) + else: + raise ValueError(f"Unrecognized message role: {message.role}") + + +class McpActorArgs(TypedDict): + name: str | None + kargs: Mapping[str, Any] + + +class McpSessionActorConfig(BaseModel): + server_params: McpServerParams + model_client: ComponentModel | Dict[str, Any] | None = None + + +class McpSessionActor(ComponentBase[BaseModel], Component[McpSessionActorConfig]): + component_type = "mcp_session_actor" + component_config_schema = McpSessionActorConfig + component_provider_override = "autogen_ext.tools.mcp.McpSessionActor" + + server_params: McpServerParams + + # model_config = ConfigDict(arbitrary_types_allowed=True) + + def __init__(self, server_params: McpServerParams, model_client: ChatCompletionClient | None = None) -> None: + self.server_params: McpServerParams = server_params + self._model_client = model_client + self.name = "mcp_session_actor" + self.description = "MCP session actor" + self._command_queue: asyncio.Queue[Dict[str, Any]] = asyncio.Queue() + self._actor_task: asyncio.Task[Any] | None = None + self._shutdown_future: asyncio.Future[Any] | None = None + self._active = False + self._initialize_result: mcp_types.InitializeResult | None = None + atexit.register(self._sync_shutdown) + + @property + def initialize_result(self) -> mcp_types.InitializeResult | None: + return self._initialize_result + + async def initialize(self) -> None: + if not self._active: + self._active = True + self._actor_task = asyncio.create_task(self._run_actor()) + + async def call(self, type: str, args: McpActorArgs | None = None) -> McpFuture: + if not self._active: + raise RuntimeError("MCP Actor not running, call initialize() first") + if self._actor_task and self._actor_task.done(): + raise RuntimeError("MCP actor task crashed", self._actor_task.exception()) + fut: asyncio.Future[McpFuture] = asyncio.Future() + if type in {"list_tools", "list_prompts", "list_resources", "list_resource_templates", "shutdown"}: + await self._command_queue.put({"type": type, "future": fut}) + res = await fut + elif type in {"call_tool", "read_resource", "get_prompt"}: + if args is None: + raise ValueError(f"args is required for {type}") + name = args.get("name", None) + kwargs = args.get("kargs", {}) + if type == "call_tool" and name is None: + raise ValueError("name is required for call_tool") + elif type == "read_resource": + uri = kwargs.get("uri", None) + if uri is None: + raise ValueError("uri is required for read_resource") + await self._command_queue.put({"type": type, "uri": uri, "future": fut}) + elif type == "get_prompt": + if name is None: + raise ValueError("name is required for get_prompt") + prompt_args = kwargs.get("arguments", None) + await self._command_queue.put({"type": type, "name": name, "args": prompt_args, "future": fut}) + else: # call_tool + await self._command_queue.put({"type": type, "name": name, "args": kwargs, "future": fut}) + res = await fut + else: + raise ValueError(f"Unknown command type: {type}") + return res + + async def close(self) -> None: + if not self._active or self._actor_task is None: + return + self._shutdown_future = asyncio.Future() + await self._command_queue.put({"type": "shutdown", "future": self._shutdown_future}) + await self._shutdown_future + await self._actor_task + self._active = False + + async def _sampling_callback( + self, + context: RequestContext[ClientSession, Any], + params: mcp_types.CreateMessageRequestParams, + ) -> mcp_types.CreateMessageResult | mcp_types.ErrorData: + """Handle sampling requests using the provided model client.""" + if self._model_client is None: + # Return an error when no model client is available + return mcp_types.ErrorData( + code=mcp_types.INVALID_REQUEST, + message="No model client available for sampling.", + data=None, + ) + + llm_messages: list[LLMMessage] = [] + + try: + if params.systemPrompt: + llm_messages.append(SystemMessage(content=params.systemPrompt)) + + for mcp_message in params.messages: + llm_messages.append(_parse_sampling_message(mcp_message, model_info=self._model_client.model_info)) + + except Exception as e: + return mcp_types.ErrorData( + code=mcp_types.INVALID_PARAMS, + message="Error processing sampling messages.", + data=f"{type(e).__name__}: {e}", + ) + + try: + result = await self._model_client.create(messages=llm_messages) + + content = result.content + if not isinstance(content, str): + content = str(content) + + return mcp_types.CreateMessageResult( + role="assistant", + content=mcp_types.TextContent(type="text", text=content), + model=self._model_client.model_info["family"], + stopReason=result.finish_reason, + ) + except Exception as e: + return mcp_types.ErrorData( + code=mcp_types.INTERNAL_ERROR, + message="Error sampling from model client.", + data=f"{type(e).__name__}: {e}", + ) + + async def _run_actor(self) -> None: + result: McpResult + try: + async with create_mcp_server_session( + self.server_params, sampling_callback=self._sampling_callback + ) as session: + # Save the initialize result + self._initialize_result = await session.initialize() + while True: + cmd = await self._command_queue.get() + if cmd["type"] == "shutdown": + cmd["future"].set_result("ok") + break + elif cmd["type"] == "call_tool": + try: + result = session.call_tool(name=cmd["name"], arguments=cmd["args"]) + cmd["future"].set_result(result) + except Exception as e: + cmd["future"].set_exception(e) + elif cmd["type"] == "read_resource": + try: + result = session.read_resource(uri=cmd["uri"]) + cmd["future"].set_result(result) + except Exception as e: + cmd["future"].set_exception(e) + elif cmd["type"] == "get_prompt": + try: + result = session.get_prompt(name=cmd["name"], arguments=cmd["args"]) + cmd["future"].set_result(result) + except Exception as e: + cmd["future"].set_exception(e) + elif cmd["type"] == "list_tools": + try: + result = session.list_tools() + cmd["future"].set_result(result) + except Exception as e: + cmd["future"].set_exception(e) + elif cmd["type"] == "list_prompts": + try: + result = session.list_prompts() + cmd["future"].set_result(result) + except Exception as e: + cmd["future"].set_exception(e) + elif cmd["type"] == "list_resources": + try: + result = session.list_resources() + cmd["future"].set_result(result) + except Exception as e: + cmd["future"].set_exception(e) + elif cmd["type"] == "list_resource_templates": + try: + result = session.list_resource_templates() + cmd["future"].set_result(result) + except Exception as e: + cmd["future"].set_exception(e) + except Exception as e: + if self._shutdown_future and not self._shutdown_future.done(): + self._shutdown_future.set_exception(e) + else: + logger.exception("Exception in MCP actor task") + finally: + self._active = False + self._actor_task = None + + def _sync_shutdown(self) -> None: + if not self._active or self._actor_task is None: + return + try: + loop = asyncio.get_event_loop() + except RuntimeError: + # No loop available — interpreter is likely shutting down + return + + if loop.is_closed(): + return + + if loop.is_running(): + loop.create_task(self.close()) + else: + loop.run_until_complete(self.close()) + + def _to_config(self) -> McpSessionActorConfig: + """ + Convert the adapter to its configuration representation. + + Returns: + McpSessionConfig: The configuration of the adapter. + """ + return McpSessionActorConfig(server_params=self.server_params) + + @classmethod + def _from_config(cls, config: McpSessionActorConfig) -> Self: + """ + Create an instance of McpSessionActor from its configuration. + + Args: + config (McpSessionConfig): The configuration of the adapter. + + Returns: + McpSessionActor: An instance of SseMcpToolAdapter. + """ + return cls(server_params=config.server_params) diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_base.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_base.py index a8bb3e939392..aabc4dcc9a95 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_base.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_base.py @@ -1,11 +1,17 @@ +import asyncio +import builtins +import json from abc import ABC -from typing import Any, Generic, Type, TypeVar +from typing import Any, Dict, Generic, Sequence, Type, TypeVar from autogen_core import CancellationToken from autogen_core.tools import BaseTool -from json_schema_to_pydantic import create_model -from mcp import Tool +from autogen_core.utils import schema_to_pydantic_model from pydantic import BaseModel +from pydantic.networks import AnyUrl + +from mcp import ClientSession, Tool +from mcp.types import AudioContent, ContentBlock, EmbeddedResource, ImageContent, ResourceLink, TextContent from ._config import McpServerParams from ._session import create_mcp_server_session @@ -24,16 +30,17 @@ class McpToolAdapter(BaseTool[BaseModel, Any], ABC, Generic[TServerParams]): component_type = "tool" - def __init__(self, server_params: TServerParams, tool: Tool) -> None: + def __init__(self, server_params: TServerParams, tool: Tool, session: ClientSession | None = None) -> None: self._tool = tool self._server_params = server_params + self._session = session # Extract name and description name = tool.name description = tool.description or "" # Create the input model from the tool's schema - input_model = create_model(tool.inputSchema) + input_model = schema_to_pydantic_model(tool.inputSchema) # Use Any as return type since MCP tool returns can vary return_type: Type[Any] = object @@ -54,22 +61,65 @@ async def run(self, args: BaseModel, cancellation_token: CancellationToken) -> A Raises: Exception: If the operation is cancelled or the tool execution fails. """ - kwargs = args.model_dump() + # Convert the input model to a dictionary + # Exclude unset values to avoid sending them to the MCP servers which may cause errors + # for many servers. + kwargs = args.model_dump(exclude_unset=True) + + if self._session is not None: + # If a session is provided, use it directly. + session = self._session + return await self._run(args=kwargs, cancellation_token=cancellation_token, session=session) + + async with create_mcp_server_session(self._server_params) as session: + await session.initialize() + return await self._run(args=kwargs, cancellation_token=cancellation_token, session=session) + + def _normalize_payload_to_content_list(self, payload: Sequence[ContentBlock]) -> list[ContentBlock]: + """ + Normalizes a raw tool output payload into a list of content items. + - If payload is already a sequence of ContentBlock items, it's converted to a list and returned. + - If payload is a single ContentBlock item, it's wrapped in a list. + - If payload is a string, it's wrapped in [TextContent(text=payload)]. + - Otherwise, the payload is stringified and wrapped in [TextContent(text=str(payload))]. + """ + if isinstance(payload, Sequence) and all( + isinstance(item, (TextContent, ImageContent, EmbeddedResource, AudioContent, ResourceLink)) + for item in payload + ): + return list(payload) + elif isinstance(payload, (TextContent, ImageContent, EmbeddedResource, AudioContent, ResourceLink)): + return [payload] + elif isinstance(payload, str): + return [TextContent(text=payload, type="text")] + else: + return [TextContent(text=str(payload), type="text")] + + async def _run(self, args: Dict[str, Any], cancellation_token: CancellationToken, session: ClientSession) -> Any: + exceptions_to_catch: tuple[Type[BaseException], ...] + if hasattr(builtins, "ExceptionGroup"): + exceptions_to_catch = (asyncio.CancelledError, builtins.ExceptionGroup) + else: + exceptions_to_catch = (asyncio.CancelledError,) try: - async with create_mcp_server_session(self._server_params) as session: - await session.initialize() + if cancellation_token.is_cancelled(): + raise asyncio.CancelledError("Operation cancelled") - if cancellation_token.is_cancelled(): - raise Exception("Operation cancelled") + result_future = asyncio.ensure_future(session.call_tool(name=self._tool.name, arguments=args)) + cancellation_token.link_future(result_future) + result = await result_future - result = await session.call_tool(self._tool.name, kwargs) # type: ignore + normalized_content_list = self._normalize_payload_to_content_list(result.content) - if result.isError: - raise Exception(f"MCP tool execution failed: {result.content}") - return result.content - except Exception as e: - raise Exception(str(e)) from e + if result.isError: + serialized_error_message = self.return_value_as_string(normalized_content_list) + raise Exception(serialized_error_message) + return normalized_content_list + + except exceptions_to_catch: + # Re-raise these specific exception types directly. + raise @classmethod async def from_server_params(cls, server_params: TServerParams, tool_name: str) -> "McpToolAdapter[TServerParams]": @@ -98,3 +148,43 @@ async def from_server_params(cls, server_params: TServerParams, tool_name: str) ) return cls(server_params=server_params, tool=matching_tool) + + def return_value_as_string(self, value: list[Any]) -> str: + """Return a string representation of the result.""" + + def serialize_item(item: Any) -> dict[str, Any]: + if isinstance(item, (TextContent, ImageContent, AudioContent)): + dumped = item.model_dump() + # Remove the 'meta' field if it exists and is None (for backward compatibility) + if dumped.get("meta") is None: + dumped.pop("meta", None) + return dumped + elif isinstance(item, EmbeddedResource): + type = item.type + resource = {} + for key, val in item.resource.model_dump().items(): + # Skip 'meta' field if it's None (for backward compatibility) + if key == "meta" and val is None: + continue + if isinstance(val, AnyUrl): + resource[key] = str(val) + else: + resource[key] = val + dumped_annotations = item.annotations.model_dump() if item.annotations else None + # Remove 'meta' from annotations if it exists and is None + if dumped_annotations and dumped_annotations.get("meta") is None: + dumped_annotations.pop("meta", None) + return {"type": type, "resource": resource, "annotations": dumped_annotations} + elif isinstance(item, ResourceLink): + dumped = item.model_dump() + # Remove the 'meta' field if it exists and is None (for backward compatibility) + if dumped.get("meta") is None: + dumped.pop("meta", None) + # Convert AnyUrl to string for JSON serialization + if "uri" in dumped and isinstance(dumped["uri"], AnyUrl): + dumped["uri"] = str(dumped["uri"]) + return dumped + else: + return {} + + return json.dumps([serialize_item(item) for item in value]) diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_config.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_config.py index 236ff6892283..d7884f489c78 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_config.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_config.py @@ -1,22 +1,42 @@ -from typing import Any, TypeAlias +from typing import Any, Literal + +from pydantic import BaseModel, Field +from typing_extensions import Annotated from mcp import StdioServerParameters -from pydantic import BaseModel class StdioServerParams(StdioServerParameters): """Parameters for connecting to an MCP server over STDIO.""" + type: Literal["StdioServerParams"] = "StdioServerParams" + read_timeout_seconds: float = 5 class SseServerParams(BaseModel): """Parameters for connecting to an MCP server over SSE.""" - url: str - headers: dict[str, Any] | None = None - timeout: float = 5 - sse_read_timeout: float = 60 * 5 + type: Literal["SseServerParams"] = "SseServerParams" + + url: str # The SSE endpoint URL. + headers: dict[str, Any] | None = None # Optional headers to include in requests. + timeout: float = 5 # HTTP timeout for regular operations. + sse_read_timeout: float = 60 * 5 # Timeout for SSE read operations. + + +class StreamableHttpServerParams(BaseModel): + """Parameters for connecting to an MCP server over Streamable HTTP.""" + + type: Literal["StreamableHttpServerParams"] = "StreamableHttpServerParams" + + url: str # The endpoint URL. + headers: dict[str, Any] | None = None # Optional headers to include in requests. + timeout: float = 30.0 # HTTP timeout for regular operations in seconds. + sse_read_timeout: float = 300.0 # Timeout for SSE read operations in seconds. + terminate_on_close: bool = True -McpServerParams: TypeAlias = StdioServerParams | SseServerParams +McpServerParams = Annotated[ + StdioServerParams | SseServerParams | StreamableHttpServerParams, Field(discriminator="type") +] diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py index 3eb8634b3698..66f8e7b7e7b3 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py @@ -1,14 +1,23 @@ -from ._config import McpServerParams, SseServerParams, StdioServerParams +from mcp import ClientSession + +from ._config import McpServerParams, SseServerParams, StdioServerParams, StreamableHttpServerParams from ._session import create_mcp_server_session from ._sse import SseMcpToolAdapter from ._stdio import StdioMcpToolAdapter +from ._streamable_http import StreamableHttpMcpToolAdapter async def mcp_server_tools( server_params: McpServerParams, -) -> list[StdioMcpToolAdapter | SseMcpToolAdapter]: + session: ClientSession | None = None, +) -> list[StdioMcpToolAdapter | SseMcpToolAdapter | StreamableHttpMcpToolAdapter]: """Creates a list of MCP tool adapters that can be used with AutoGen agents. + .. warning:: + + Only connect to trusted MCP servers, especially when using + `StdioServerParams` as it executes commands in the local environment. + This factory function connects to an MCP server and returns adapters for all available tools. The adapters can be directly assigned to an AutoGen agent's tools list. @@ -23,11 +32,14 @@ async def mcp_server_tools( Args: server_params (McpServerParams): Connection parameters for the MCP server. Can be either StdioServerParams for command-line tools or - SseServerParams for HTTP/SSE services. + SseServerParams and StreamableHttpServerParams for HTTP/SSE services. + session (ClientSession | None): Optional existing session to use. This is used + when you want to reuse an existing connection to the MCP server. The session + will be reused when creating the MCP tool adapters. Returns: - list[StdioMcpToolAdapter | SseMcpToolAdapter]: A list of tool adapters ready to use - with AutoGen agents. + list[StdioMcpToolAdapter | SseMcpToolAdapter | StreamableHttpMcpToolAdapter]: + A list of tool adapters ready to use with AutoGen agents. Examples: @@ -105,7 +117,59 @@ async def main() -> None: # Let the agent fetch the content of a URL and summarize it. result = await agent.run(task="Summarize the content of https://en.wikipedia.org/wiki/Seattle") - print(result.messages[-1].content) + print(result.messages[-1]) + + + asyncio.run(main()) + + **Sharing an MCP client session across multiple tools:** + + You can create a single MCP client session and share it across multiple tools. + This is sometimes required when the server maintains a session state + (e.g., a browser state) that should be reused for multiple requests. + + The following example show how to create a single MCP client session + to a local `Playwright `_ + server and share it across multiple tools. + + + .. code-block:: python + + import asyncio + + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.conditions import TextMentionTermination + from autogen_agentchat.teams import RoundRobinGroupChat + from autogen_agentchat.ui import Console + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_ext.tools.mcp import StdioServerParams, create_mcp_server_session, mcp_server_tools + + + async def main() -> None: + model_client = OpenAIChatCompletionClient(model="gpt-4o", parallel_tool_calls=False) # type: ignore + params = StdioServerParams( + command="npx", + args=["@playwright/mcp@latest"], + read_timeout_seconds=60, + ) + async with create_mcp_server_session(params) as session: + await session.initialize() + tools = await mcp_server_tools(server_params=params, session=session) + print(f"Tools: {[tool.name for tool in tools]}") + + agent = AssistantAgent( + name="Assistant", + model_client=model_client, + tools=tools, # type: ignore + ) + + termination = TextMentionTermination("TERMINATE") + team = RoundRobinGroupChat([agent], termination_condition=termination) + await Console( + team.run_stream( + task="Go to https://ekzhu.com/, visit the first link in the page, then tell me about the linked page." + ) + ) asyncio.run(main()) @@ -130,13 +194,21 @@ async def main() -> None: For more examples and detailed usage, see the samples directory in the package repository. """ - async with create_mcp_server_session(server_params) as session: - await session.initialize() + if session is None: + async with create_mcp_server_session(server_params) as temp_session: + await temp_session.initialize() + tools = await temp_session.list_tools() + else: tools = await session.list_tools() if isinstance(server_params, StdioServerParams): - return [StdioMcpToolAdapter(server_params=server_params, tool=tool) for tool in tools.tools] + return [StdioMcpToolAdapter(server_params=server_params, tool=tool, session=session) for tool in tools.tools] elif isinstance(server_params, SseServerParams): - return [SseMcpToolAdapter(server_params=server_params, tool=tool) for tool in tools.tools] + return [SseMcpToolAdapter(server_params=server_params, tool=tool, session=session) for tool in tools.tools] + elif isinstance(server_params, StreamableHttpServerParams): + return [ + StreamableHttpMcpToolAdapter(server_params=server_params, tool=tool, session=session) + for tool in tools.tools + ] raise ValueError(f"Unsupported server params type: {type(server_params)}") diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_session.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_session.py index bc1a28fac5cb..04ea17fedef6 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_session.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_session.py @@ -3,15 +3,17 @@ from typing import AsyncGenerator from mcp import ClientSession +from mcp.client.session import SamplingFnT from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client +from mcp.client.streamable_http import streamablehttp_client -from ._config import McpServerParams, SseServerParams, StdioServerParams +from ._config import McpServerParams, SseServerParams, StdioServerParams, StreamableHttpServerParams @asynccontextmanager async def create_mcp_server_session( - server_params: McpServerParams, + server_params: McpServerParams, sampling_callback: SamplingFnT | None = None ) -> AsyncGenerator[ClientSession, None]: """Create an MCP client session for the given server parameters.""" if isinstance(server_params, StdioServerParams): @@ -20,9 +22,34 @@ async def create_mcp_server_session( read_stream=read, write_stream=write, read_timeout_seconds=timedelta(seconds=server_params.read_timeout_seconds), + sampling_callback=sampling_callback, ) as session: yield session elif isinstance(server_params, SseServerParams): - async with sse_client(**server_params.model_dump()) as (read, write): - async with ClientSession(read_stream=read, write_stream=write) as session: + async with sse_client(**server_params.model_dump(exclude={"type"})) as (read, write): + async with ClientSession( + read_stream=read, + write_stream=write, + read_timeout_seconds=timedelta(seconds=server_params.sse_read_timeout), + sampling_callback=sampling_callback, + ) as session: + yield session + elif isinstance(server_params, StreamableHttpServerParams): + # Convert float seconds to timedelta for the streamablehttp_client + params_dict = server_params.model_dump(exclude={"type"}) + params_dict["timeout"] = timedelta(seconds=server_params.timeout) + params_dict["sse_read_timeout"] = timedelta(seconds=server_params.sse_read_timeout) + + async with streamablehttp_client(**params_dict) as ( + read, + write, + session_id_callback, # type: ignore[assignment, unused-variable] + ): + # TODO: Handle session_id_callback if needed + async with ClientSession( + read_stream=read, + write_stream=write, + read_timeout_seconds=timedelta(seconds=server_params.sse_read_timeout), + sampling_callback=sampling_callback, + ) as session: yield session diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_sse.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_sse.py index 252af7ce50da..c77ec8607422 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_sse.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_sse.py @@ -1,8 +1,9 @@ from autogen_core import Component -from mcp import Tool from pydantic import BaseModel from typing_extensions import Self +from mcp import ClientSession, Tool + from ._base import McpToolAdapter from ._config import SseServerParams @@ -35,8 +36,11 @@ class SseMcpToolAdapter( Args: server_params (SseServerParameters): Parameters for the MCP server connection, - including URL, headers, and timeouts - tool (Tool): The MCP tool to wrap + including URL, headers, and timeouts. + tool (Tool): The MCP tool to wrap. + session (ClientSession, optional): The MCP client session to use. If not provided, + it will create a new session. This is useful for testing or when you want to + manage the session lifecycle yourself. Examples: Use a remote translation service that implements MCP over SSE to create tools @@ -86,8 +90,8 @@ async def main() -> None: component_config_schema = SseMcpToolAdapterConfig component_provider_override = "autogen_ext.tools.mcp.SseMcpToolAdapter" - def __init__(self, server_params: SseServerParams, tool: Tool) -> None: - super().__init__(server_params=server_params, tool=tool) + def __init__(self, server_params: SseServerParams, tool: Tool, session: ClientSession | None = None) -> None: + super().__init__(server_params=server_params, tool=tool, session=session) def _to_config(self) -> SseMcpToolAdapterConfig: """ diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_stdio.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_stdio.py index 4f827785e903..bbe7c6ca0752 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_stdio.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_stdio.py @@ -1,8 +1,9 @@ from autogen_core import Component -from mcp import Tool from pydantic import BaseModel from typing_extensions import Self +from mcp import ClientSession, Tool + from ._base import McpToolAdapter from ._config import StdioServerParams @@ -37,6 +38,9 @@ class StdioMcpToolAdapter( server_params (StdioServerParams): Parameters for the MCP server connection, including command to run and its arguments tool (Tool): The MCP tool to wrap + session (ClientSession, optional): The MCP client session to use. If not provided, + a new session will be created. This is useful for testing or when you want to + manage the session lifecycle yourself. See :func:`~autogen_ext.tools.mcp.mcp_server_tools` for examples. """ @@ -44,8 +48,8 @@ class StdioMcpToolAdapter( component_config_schema = StdioMcpToolAdapterConfig component_provider_override = "autogen_ext.tools.mcp.StdioMcpToolAdapter" - def __init__(self, server_params: StdioServerParams, tool: Tool) -> None: - super().__init__(server_params=server_params, tool=tool) + def __init__(self, server_params: StdioServerParams, tool: Tool, session: ClientSession | None = None) -> None: + super().__init__(server_params=server_params, tool=tool, session=session) def _to_config(self) -> StdioMcpToolAdapterConfig: """ diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_streamable_http.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_streamable_http.py new file mode 100644 index 000000000000..a7df719c2703 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_streamable_http.py @@ -0,0 +1,121 @@ +from autogen_core import Component +from pydantic import BaseModel +from typing_extensions import Self + +from mcp import ClientSession, Tool + +from ._base import McpToolAdapter +from ._config import StreamableHttpServerParams + + +class StreamableHttpMcpToolAdapterConfig(BaseModel): + """Configuration for the MCP tool adapter.""" + + server_params: StreamableHttpServerParams + tool: Tool + + +class StreamableHttpMcpToolAdapter( + McpToolAdapter[StreamableHttpServerParams], + Component[StreamableHttpMcpToolAdapterConfig], +): + """ + Allows you to wrap an MCP tool running over Streamable HTTP and make it available to AutoGen. + + This adapter enables using MCP-compatible tools that communicate over Streamable HTTP + with AutoGen agents. Common use cases include integrating with remote MCP services, + cloud-based tools, and web APIs that implement the Model Context Protocol (MCP). + + .. note:: + + To use this class, you need to install `mcp` extra for the `autogen-ext` package. + + .. code-block:: bash + + pip install -U "autogen-ext[mcp]" + + + Args: + server_params (StreamableHttpServerParams): Parameters for the MCP server connection, + including URL, headers, and timeouts. + tool (Tool): The MCP tool to wrap. + session (ClientSession, optional): The MCP client session to use. If not provided, + it will create a new session. This is useful for testing or when you want to + manage the session lifecycle yourself. + + Examples: + Use a remote translation service that implements MCP over Streamable HTTP to + create tools that allow AutoGen agents to perform translations: + + .. code-block:: python + + import asyncio + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_ext.tools.mcp import StreamableHttpMcpToolAdapter, StreamableHttpServerParams + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.ui import Console + from autogen_core import CancellationToken + + + async def main() -> None: + # Create server params for the remote MCP service + server_params = StreamableHttpServerParams( + url="https://api.example.com/mcp", + headers={"Authorization": "Bearer your-api-key", "Content-Type": "application/json"}, + timeout=30.0, # HTTP timeout in seconds + sse_read_timeout=300.0, # SSE read timeout in seconds (5 minutes) + terminate_on_close=True, + ) + + # Get the translation tool from the server + adapter = await StreamableHttpMcpToolAdapter.from_server_params(server_params, "translate") + + # Create an agent that can use the translation tool + model_client = OpenAIChatCompletionClient(model="gpt-4") + agent = AssistantAgent( + name="translator", + model_client=model_client, + tools=[adapter], + system_message="You are a helpful translation assistant.", + ) + + # Let the agent translate some text + await Console( + agent.run_stream(task="Translate 'Hello, how are you?' to Spanish", cancellation_token=CancellationToken()) + ) + + + if __name__ == "__main__": + asyncio.run(main()) + + """ + + component_config_schema = StreamableHttpMcpToolAdapterConfig + component_provider_override = "autogen_ext.tools.mcp.StreamableHttpMcpToolAdapter" + + def __init__( + self, server_params: StreamableHttpServerParams, tool: Tool, session: ClientSession | None = None + ) -> None: + super().__init__(server_params=server_params, tool=tool, session=session) + + def _to_config(self) -> StreamableHttpMcpToolAdapterConfig: + """ + Convert the adapter to its configuration representation. + + Returns: + StreamableHttpMcpToolAdapterConfig: The configuration of the adapter. + """ + return StreamableHttpMcpToolAdapterConfig(server_params=self._server_params, tool=self._tool) + + @classmethod + def _from_config(cls, config: StreamableHttpMcpToolAdapterConfig) -> Self: + """ + Create an instance of StreamableHttpMcpToolAdapter from its configuration. + + Args: + config (StreamableHttpMcpToolAdapterConfig): The configuration of the adapter. + + Returns: + StreamableHttpMcpToolAdapter: An instance of StreamableHttpMcpToolAdapter. + """ + return cls(server_params=config.server_params, tool=config.tool) diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py new file mode 100644 index 000000000000..ff12f0d072c9 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py @@ -0,0 +1,518 @@ +import asyncio +import builtins +import warnings +from typing import Any, Dict, List, Literal, Mapping, Optional + +from autogen_core import CancellationToken, Component, ComponentModel, Image, trace_tool_span +from autogen_core.models import ChatCompletionClient +from autogen_core.tools import ( + ImageResultContent, + ParametersSchema, + TextResultContent, + ToolOverride, + ToolResult, + ToolSchema, + Workbench, +) +from pydantic import BaseModel, Field +from typing_extensions import Self + +from mcp.types import ( + CallToolResult, + EmbeddedResource, + GetPromptResult, + ImageContent, + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, + ListToolsResult, + ReadResourceResult, + TextContent, +) + +from ._actor import McpSessionActor +from ._config import McpServerParams, SseServerParams, StdioServerParams, StreamableHttpServerParams + + +class McpWorkbenchConfig(BaseModel): + server_params: McpServerParams + tool_overrides: Dict[str, ToolOverride] = Field(default_factory=dict) + model_client: ComponentModel | Dict[str, Any] | None = None + + +class McpWorkbenchState(BaseModel): + type: Literal["McpWorkBenchState"] = "McpWorkBenchState" + + +class McpWorkbench(Workbench, Component[McpWorkbenchConfig]): + """A workbench that wraps an MCP server and provides an interface + to list and call tools provided by the server. + + .. warning:: + + Only connect to trusted MCP servers, especially when using + `StdioServerParams` as it executes commands in the local environment. + + This workbench should be used as a context manager to ensure proper + initialization and cleanup of the underlying MCP session. + + .. list-table:: MCP Support + :header-rows: 1 + :widths: 30 70 + + * - MCP Capability + - Supported Features + * - Tools + - list_tools, call_tool + * - Resources + - list_resources, read_resource + * - ResourceTemplates + - list_resource_templates, read_resource_template + * - Prompts + - list_prompts, get_prompt + * - Sampling + - Optional support via model_client + * - Roots + - not supported + * - Ellicitation + - not supported + + Args: + server_params (McpServerParams): The parameters to connect to the MCP server. + This can be either a :class:`StdioServerParams` or :class:`SseServerParams`. + tool_overrides (Optional[Dict[str, ToolOverride]]): Optional mapping of original tool + names to override configurations for name and/or description. This allows + customizing how server tools appear to consumers while maintaining the underlying + tool functionality. + model_client: Optional chat completion client to handle sampling requests + from MCP servers that support the sampling capability. This allows MCP + servers to request text generation from a language model during tool + execution. If not provided, sampling requests will return an error. + + Raises: + ValueError: If there are conflicts in tool override names. + + Examples: + + Here is a simple example of how to use the workbench with a `mcp-server-fetch` server: + + .. code-block:: python + + import asyncio + + from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams + + + async def main() -> None: + params = StdioServerParams( + command="uvx", + args=["mcp-server-fetch"], + read_timeout_seconds=60, + ) + + # You can also use `start()` and `stop()` to manage the session. + async with McpWorkbench(server_params=params) as workbench: + tools = await workbench.list_tools() + print(tools) + result = await workbench.call_tool(tools[0]["name"], {"url": "https://github.com/"}) + print(result) + + + asyncio.run(main()) + + Example of using tool overrides: + + .. code-block:: python + + import asyncio + from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams + from autogen_core.tools import ToolOverride + + + async def main() -> None: + params = StdioServerParams( + command="uvx", + args=["mcp-server-fetch"], + read_timeout_seconds=60, + ) + + # Override the fetch tool's name and description + overrides = { + "fetch": ToolOverride(name="web_fetch", description="Enhanced web fetching tool with better error handling") + } + + async with McpWorkbench(server_params=params, tool_overrides=overrides) as workbench: + tools = await workbench.list_tools() + # The tool will now appear as "web_fetch" with the new description + print(tools) + # Call the overridden tool + result = await workbench.call_tool("web_fetch", {"url": "https://github.com/"}) + print(result) + + + asyncio.run(main()) + + Example of using the workbench with the `GitHub MCP Server `_: + + .. code-block:: python + + import asyncio + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.ui import Console + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams + + + async def main() -> None: + model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano") + server_params = StdioServerParams( + command="docker", + args=[ + "run", + "-i", + "--rm", + "-e", + "GITHUB_PERSONAL_ACCESS_TOKEN", + "ghcr.io/github/github-mcp-server", + ], + env={ + "GITHUB_PERSONAL_ACCESS_TOKEN": "ghp_XXXXXXXXXXXXXXXXXXXXXXXXXXXXXX", + }, + ) + async with McpWorkbench(server_params) as mcp: + agent = AssistantAgent( + "github_assistant", + model_client=model_client, + workbench=mcp, + reflect_on_tool_use=True, + model_client_stream=True, + ) + await Console(agent.run_stream(task="Is there a repository named Autogen")) + + + asyncio.run(main()) + + Example of using the workbench with the `Playwright MCP Server `_: + + .. code-block:: python + + # First run `npm install -g @playwright/mcp@latest` to install the MCP server. + import asyncio + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.teams import RoundRobinGroupChat + from autogen_agentchat.conditions import TextMessageTermination + from autogen_agentchat.ui import Console + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams + + + async def main() -> None: + model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano") + server_params = StdioServerParams( + command="npx", + args=[ + "@playwright/mcp@latest", + "--headless", + ], + ) + async with McpWorkbench(server_params) as mcp: + agent = AssistantAgent( + "web_browsing_assistant", + model_client=model_client, + workbench=mcp, + model_client_stream=True, + ) + team = RoundRobinGroupChat( + [agent], + termination_condition=TextMessageTermination(source="web_browsing_assistant"), + ) + await Console(team.run_stream(task="Find out how many contributors for the microsoft/autogen repository")) + + + asyncio.run(main()) + + """ + + component_provider_override = "autogen_ext.tools.mcp.McpWorkbench" + component_config_schema = McpWorkbenchConfig + + def __init__( + self, + server_params: McpServerParams, + tool_overrides: Optional[Dict[str, ToolOverride]] = None, + model_client: ChatCompletionClient | None = None, + ) -> None: + self._server_params = server_params + self._tool_overrides = tool_overrides or {} + self._model_client = model_client + + # Build reverse mapping from override names to original names for call_tool + self._override_name_to_original: Dict[str, str] = {} + for original_name, override in self._tool_overrides.items(): + override_name = override.name + if override_name and override_name != original_name: + # Check for conflicts with other override names + if override_name in self._override_name_to_original: + existing_original = self._override_name_to_original[override_name] + raise ValueError( + f"Tool override name '{override_name}' is used by multiple tools: " + f"'{existing_original}' and '{original_name}'. Override names must be unique." + ) + self._override_name_to_original[override_name] = original_name + + # self._session: ClientSession | None = None + self._actor: McpSessionActor | None = None + self._actor_loop: asyncio.AbstractEventLoop | None = None + self._read = None + self._write = None + + @property + def server_params(self) -> McpServerParams: + return self._server_params + + async def list_tools(self) -> List[ToolSchema]: + if not self._actor: + await self.start() # fallback to start the actor if not initialized instead of raising an error + # Why? Because when deserializing the workbench, the actor might not be initialized yet. + # raise RuntimeError("Actor is not initialized. Call start() first.") + if self._actor is None: + raise RuntimeError("Actor is not initialized. Please check the server connection.") + result_future = await self._actor.call("list_tools", None) + list_tool_result = await result_future + assert isinstance( + list_tool_result, ListToolsResult + ), f"list_tools must return a CallToolResult, instead of : {str(type(list_tool_result))}" + schema: List[ToolSchema] = [] + for tool in list_tool_result.tools: + original_name = tool.name + name = original_name + description = tool.description or "" + + # Apply overrides if they exist for this tool + if original_name in self._tool_overrides: + override = self._tool_overrides[original_name] + if override.name is not None: + name = override.name + if override.description is not None: + description = override.description + + parameters = ParametersSchema( + type="object", + properties=tool.inputSchema.get("properties", {}), + required=tool.inputSchema.get("required", []), + additionalProperties=tool.inputSchema.get("additionalProperties", False), + ) + tool_schema = ToolSchema( + name=name, + description=description, + parameters=parameters, + ) + schema.append(tool_schema) + return schema + + async def call_tool( + self, + name: str, + arguments: Mapping[str, Any] | None = None, + cancellation_token: CancellationToken | None = None, + call_id: str | None = None, + ) -> ToolResult: + if not self._actor: + await self.start() # fallback to start the actor if not initialized instead of raising an error + # Why? Because when deserializing the workbench, the actor might not be initialized yet. + # raise RuntimeError("Actor is not initialized. Call start() first.") + if self._actor is None: + raise RuntimeError("Actor is not initialized. Please check the server connection.") + if not cancellation_token: + cancellation_token = CancellationToken() + if not arguments: + arguments = {} + + # Check if the name is an override name and map it back to the original + original_name = self._override_name_to_original.get(name, name) + + with trace_tool_span( + tool_name=name, # Use the requested name for tracing + tool_call_id=call_id, + ): + try: + result_future = await self._actor.call("call_tool", {"name": original_name, "kargs": arguments}) + cancellation_token.link_future(result_future) + result = await result_future + assert isinstance( + result, CallToolResult + ), f"call_tool must return a CallToolResult, instead of : {str(type(result))}" + result_parts: List[TextResultContent | ImageResultContent] = [] + is_error = result.isError + for content in result.content: + if isinstance(content, TextContent): + result_parts.append(TextResultContent(content=content.text)) + elif isinstance(content, ImageContent): + result_parts.append(ImageResultContent(content=Image.from_base64(content.data))) + elif isinstance(content, EmbeddedResource): + # TODO: how to handle embedded resources? + # For now we just use text representation. + result_parts.append(TextResultContent(content=content.model_dump_json())) + else: + raise ValueError(f"Unknown content type from server: {type(content)}") + except Exception as e: + error_message = self._format_errors(e) + is_error = True + result_parts = [TextResultContent(content=error_message)] + return ToolResult(name=name, result=result_parts, is_error=is_error) # Return the requested name + + @property + def initialize_result(self) -> Any: + if self._actor: + return self._actor.initialize_result + + return None + + async def list_prompts(self) -> ListPromptsResult: + """List available prompts from the MCP server.""" + if not self._actor: + await self.start() + if self._actor is None: + raise RuntimeError("Actor is not initialized. Please check the server connection.") + + result_future = await self._actor.call("list_prompts", None) + list_prompts_result = await result_future + assert isinstance( + list_prompts_result, ListPromptsResult + ), f"list_prompts must return a ListPromptsResult, instead of: {str(type(list_prompts_result))}" + + return list_prompts_result + + async def list_resources(self) -> ListResourcesResult: + """List available resources from the MCP server.""" + if not self._actor: + await self.start() + if self._actor is None: + raise RuntimeError("Actor is not initialized. Please check the server connection.") + + result_future = await self._actor.call("list_resources", None) + list_resources_result = await result_future + assert isinstance( + list_resources_result, ListResourcesResult + ), f"list_resources must return a ListResourcesResult, instead of: {str(type(list_resources_result))}" + + return list_resources_result + + async def list_resource_templates(self) -> ListResourceTemplatesResult: + """List available resource templates from the MCP server.""" + if not self._actor: + await self.start() + if self._actor is None: + raise RuntimeError("Actor is not initialized. Please check the server connection.") + + result_future = await self._actor.call("list_resource_templates", None) + list_templates_result = await result_future + assert isinstance( + list_templates_result, ListResourceTemplatesResult + ), f"list_resource_templates must return a ListResourceTemplatesResult, instead of: {str(type(list_templates_result))}" + + return list_templates_result + + async def read_resource(self, uri: str) -> ReadResourceResult: + """Read a resource from the MCP server.""" + if not self._actor: + await self.start() + if self._actor is None: + raise RuntimeError("Actor is not initialized. Please check the server connection.") + + result_future = await self._actor.call("read_resource", {"name": None, "kargs": {"uri": uri}}) + read_resource_result = await result_future + assert isinstance( + read_resource_result, ReadResourceResult + ), f"read_resource must return a ReadResourceResult, instead of: {str(type(read_resource_result))}" + + return read_resource_result + + async def get_prompt(self, name: str, arguments: Optional[Dict[str, str]] = None) -> GetPromptResult: + """Get a prompt from the MCP server.""" + if not self._actor: + await self.start() + if self._actor is None: + raise RuntimeError("Actor is not initialized. Please check the server connection.") + + result_future = await self._actor.call("get_prompt", {"name": name, "kargs": {"arguments": arguments}}) + get_prompt_result = await result_future + assert isinstance( + get_prompt_result, GetPromptResult + ), f"get_prompt must return a GetPromptResult, instead of: {str(type(get_prompt_result))}" + + return get_prompt_result + + def _format_errors(self, error: Exception) -> str: + """Recursively format errors into a string.""" + + error_message = "" + if hasattr(builtins, "ExceptionGroup") and isinstance(error, builtins.ExceptionGroup): + # ExceptionGroup is available in Python 3.11+. + # TODO: how to make this compatible with Python 3.10? + for sub_exception in error.exceptions: # type: ignore + error_message += self._format_errors(sub_exception) # type: ignore + else: + error_message += f"{str(error)}\n" + return error_message + + async def start(self) -> None: + if self._actor: + warnings.warn( + "McpWorkbench is already started. No need to start again.", + UserWarning, + stacklevel=2, + ) + return # Already initialized, no need to start again + + if isinstance(self._server_params, (StdioServerParams, SseServerParams, StreamableHttpServerParams)): + self._actor = McpSessionActor(self._server_params, model_client=self._model_client) + await self._actor.initialize() + self._actor_loop = asyncio.get_event_loop() + else: + raise ValueError(f"Unsupported server params type: {type(self._server_params)}") + + async def stop(self) -> None: + if self._actor: + # Close the actor + await self._actor.close() + self._actor = None + else: + raise RuntimeError("McpWorkbench is not started. Call start() first.") + + async def reset(self) -> None: + pass + + async def save_state(self) -> Mapping[str, Any]: + return McpWorkbenchState().model_dump() + + async def load_state(self, state: Mapping[str, Any]) -> None: + pass + + def _to_config(self) -> McpWorkbenchConfig: + model_client_config = None + if self._model_client is not None: + model_client_config = self._model_client.dump_component() + return McpWorkbenchConfig( + server_params=self._server_params, tool_overrides=self._tool_overrides, model_client=model_client_config + ) + + @classmethod + def _from_config(cls, config: McpWorkbenchConfig) -> Self: + model_client = None + if config.model_client is not None: + model_client = ChatCompletionClient.load_component(config.model_client) + return cls(server_params=config.server_params, tool_overrides=config.tool_overrides, model_client=model_client) + + def __del__(self) -> None: + # Ensure the actor is stopped when the workbench is deleted + # Use getattr to safely handle cases where attributes may not be set (e.g., if __init__ failed) + actor = getattr(self, "_actor", None) + actor_loop = getattr(self, "_actor_loop", None) + + if actor and actor_loop: + if actor_loop.is_running() and not actor_loop.is_closed(): + actor_loop.call_soon_threadsafe(lambda: asyncio.create_task(self.stop())) + else: + msg = "Cannot safely stop actor at [McpWorkbench.__del__]: loop is closed or not running" + warnings.warn(msg, RuntimeWarning, stacklevel=2) diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/semantic_kernel/__init__.py b/python/packages/autogen-ext/src/autogen_ext/tools/semantic_kernel/__init__.py index eed4547717da..358a71743809 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/semantic_kernel/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/semantic_kernel/__init__.py @@ -1,5 +1,6 @@ -from ._kernel_function_from_tool import KernelFunctionFromTool +from ._kernel_function_from_tool import KernelFunctionFromTool, KernelFunctionFromToolSchema __all__ = [ "KernelFunctionFromTool", + "KernelFunctionFromToolSchema", ] diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/semantic_kernel/_kernel_function_from_tool.py b/python/packages/autogen-ext/src/autogen_ext/tools/semantic_kernel/_kernel_function_from_tool.py index 0f3f93e2e7fa..114919558b20 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/semantic_kernel/_kernel_function_from_tool.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/semantic_kernel/_kernel_function_from_tool.py @@ -1,10 +1,13 @@ from typing import Any, TypeVar from autogen_core import CancellationToken -from autogen_core.tools import BaseTool +from autogen_core.tools import BaseTool, ToolSchema from pydantic import BaseModel -from semantic_kernel.functions import KernelFunctionFromMethod, kernel_function + +from semantic_kernel.functions import KernelFunctionFromMethod, KernelFunctionFromPrompt, kernel_function from semantic_kernel.functions.kernel_parameter_metadata import KernelParameterMetadata +from semantic_kernel.prompt_template.input_variable import InputVariable +from semantic_kernel.prompt_template.prompt_template_config import PromptTemplateConfig InputT = TypeVar("InputT", bound=BaseModel) OutputT = TypeVar("OutputT", bound=BaseModel) @@ -65,3 +68,27 @@ async def tool_method(**kwargs: dict[str, Any]) -> Any: ) self._tool = tool + + +class KernelFunctionFromToolSchema(KernelFunctionFromPrompt): + def __init__(self, tool_schema: ToolSchema, plugin_name: str | None = None): + properties = tool_schema.get("parameters", {}).get("properties", {}) + required = properties.get("required", []) + + prompt_template_config = PromptTemplateConfig( + name=tool_schema.get("name", ""), + description=tool_schema.get("description", ""), + input_variables=[ + InputVariable( + name=prop_name, description=prop_info.get("description", ""), is_required=prop_name in required + ) + for prop_name, prop_info in properties.items() + ], + ) + + super().__init__( + function_name=tool_schema.get("name", ""), + plugin_name=plugin_name, + description=tool_schema.get("description", ""), + prompt_template_config=prompt_template_config, + ) diff --git a/python/packages/autogen-ext/src/autogen_ext/ui/_rich_console.py b/python/packages/autogen-ext/src/autogen_ext/ui/_rich_console.py index 1951205e8ed5..04299ab06824 100644 --- a/python/packages/autogen-ext/src/autogen_ext/ui/_rich_console.py +++ b/python/packages/autogen-ext/src/autogen_ext/ui/_rich_console.py @@ -14,8 +14,8 @@ from autogen_agentchat.base import Response, TaskResult from autogen_agentchat.messages import ( - AgentEvent, - ChatMessage, + BaseAgentEvent, + BaseChatMessage, ModelClientStreamingChunkEvent, MultiModalMessage, UserInputRequestedEvent, @@ -56,12 +56,12 @@ def aprint(output: str, end: str = "\n") -> Awaitable[None]: return asyncio.to_thread(print, output, end=end) -def _extract_message_content(message: AgentEvent | ChatMessage) -> Tuple[List[str], List[Image]]: +def _extract_message_content(message: BaseAgentEvent | BaseChatMessage) -> Tuple[List[str], List[Image]]: if isinstance(message, MultiModalMessage): text_parts = [item for item in message.content if isinstance(item, str)] image_parts = [item for item in message.content if isinstance(item, Image)] else: - text_parts = [str(message.content)] + text_parts = [message.to_text()] image_parts = [] return text_parts, image_parts @@ -100,7 +100,7 @@ async def _aprint_message_content( async def RichConsole( - stream: AsyncGenerator[AgentEvent | ChatMessage | T, None], + stream: AsyncGenerator[BaseAgentEvent | BaseChatMessage | T, None], *, no_inline_images: bool = False, output_stats: bool = False, @@ -117,7 +117,7 @@ async def RichConsole( It will be improved in future releases. Args: - stream (AsyncGenerator[AgentEvent | ChatMessage | TaskResult, None] | AsyncGenerator[AgentEvent | ChatMessage | Response, None]): Message stream to render. + stream (AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None] | AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]): Message stream to render. This can be from :meth:`~autogen_agentchat.base.TaskRunner.run_stream` or :meth:`~autogen_agentchat.base.ChatAgent.on_messages_stream`. no_inline_images (bool, optional): If terminal is iTerm2 will render images inline. Use this to disable this behavior. Defaults to False. output_stats (bool, optional): (Experimental) If True, will output a summary of the messages and inline token usage info. Defaults to False. @@ -191,7 +191,7 @@ async def RichConsole( pass else: # Cast required for mypy to be happy - message = cast(AgentEvent | ChatMessage, message) # type: ignore + message = cast(BaseAgentEvent | BaseChatMessage, message) # type: ignore text_parts, image_parts = _extract_message_content(message) # Add usage stats if needed diff --git a/python/packages/autogen-ext/tests/agents/test_openai_agent_builtin_tool_validation.py b/python/packages/autogen-ext/tests/agents/test_openai_agent_builtin_tool_validation.py new file mode 100644 index 000000000000..c66df4bb2350 --- /dev/null +++ b/python/packages/autogen-ext/tests/agents/test_openai_agent_builtin_tool_validation.py @@ -0,0 +1,754 @@ +"""Tests for OpenAI agent builtin tool validation.""" + +# Standard library imports +import os +from typing import Any, Dict, cast + +# Third-party imports +import pytest + +# Local imports +from autogen_agentchat.messages import TextMessage +from autogen_core import CancellationToken +from autogen_ext.agents.openai import OpenAIAgent +from openai import AsyncOpenAI +from pytest import MonkeyPatch + + +@pytest.fixture(autouse=True) +def set_dummy_openai_key(monkeypatch: MonkeyPatch) -> None: + """Ensure tests have a dummy OPENAI_API_KEY by default.""" + # Only set a dummy key if no api key is provided + if not os.getenv("OPENAI_API_KEY"): + monkeypatch.setenv("OPENAI_API_KEY", "sk-test-dummy-key") + + +skip_if_no_real_openai_key = pytest.mark.skipif( + os.getenv("OPENAI_API_KEY", "") in ("", "sk-test-dummy-key"), + reason="No real OPENAI_API_KEY provided; skipping integration test.", +) + + +@pytest.fixture +def openai_client() -> AsyncOpenAI: + """Provides an AsyncOpenAI client using the test API key.""" + return AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY", "")) + + +@pytest.fixture +def cancel_token() -> CancellationToken: + """Provides a fresh CancellationToken for each test.""" + return CancellationToken() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_name,model,should_raise", + [ + ("web_search_preview", "gpt-4o", False), + ("image_generation", "gpt-4o", False), + ("local_shell", "codex-mini-latest", False), + ("local_shell", "gpt-4o", True), + ("file_search", "gpt-4o", True), + ("code_interpreter", "gpt-4o", True), + ("computer_use_preview", "gpt-4o", True), + ("mcp", "gpt-4o", True), + ("not_a_tool", "gpt-4o", True), + ], +) +async def test_builtin_tool_string_validation( + tool_name: str, model: str, should_raise: bool, openai_client: AsyncOpenAI +) -> None: + """Test validation of string-based builtin tools.""" + client = openai_client + tools = [tool_name] # type: ignore + + if should_raise: + with pytest.raises(ValueError): + OpenAIAgent( + name="test", + description="desc", + client=client, + model=model, + instructions="inst", + tools=tools, # type: ignore + ) + else: + agent = OpenAIAgent( + name="test", + description="desc", + client=client, + model=model, + instructions="inst", + tools=tools, # type: ignore + ) + assert any(t["type"] == tool_name for t in agent.tools) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_config,should_raise", + [ + # file_search: missing required param + ({"type": "file_search"}, True), + # file_search: empty vector_store_ids + ({"type": "file_search", "vector_store_ids": []}, True), + # file_search: invalid type + ({"type": "file_search", "vector_store_ids": [123]}, True), + # file_search: valid + ({"type": "file_search", "vector_store_ids": ["vs1"]}, False), + # computer_use_preview: missing param + ({"type": "computer_use_preview", "display_height": 100, "display_width": 100}, True), + # computer_use_preview: invalid type + ({"type": "computer_use_preview", "display_height": -1, "display_width": 100, "environment": "desktop"}, True), + # computer_use_preview: valid + ( + {"type": "computer_use_preview", "display_height": 100, "display_width": 100, "environment": "desktop"}, + False, + ), + # code_interpreter: missing param + ({"type": "code_interpreter"}, True), + # code_interpreter: empty container + ({"type": "code_interpreter", "container": ""}, True), + # code_interpreter: valid + ({"type": "code_interpreter", "container": "python-3.11"}, False), + # mcp: missing param + ({"type": "mcp", "server_label": "label"}, True), + # mcp: invalid type + ({"type": "mcp", "server_label": "", "server_url": "url"}, True), + # mcp: valid + ({"type": "mcp", "server_label": "label", "server_url": "url"}, False), + # web_search_preview: valid with string user_location + ({"type": "web_search_preview", "user_location": "US"}, False), + # web_search_preview: valid with dict user_location + ({"type": "web_search_preview", "user_location": {"type": "approximate"}}, False), + # web_search_preview: invalid user_location type + ({"type": "web_search_preview", "user_location": 123}, True), + # image_generation: valid with background + ({"type": "image_generation", "background": "white"}, False), + # image_generation: invalid background + ({"type": "image_generation", "background": ""}, True), + ], +) +async def test_builtin_tool_dict_validation( + tool_config: Dict[str, Any], should_raise: bool, openai_client: AsyncOpenAI +) -> None: + """Test validation of dictionary-based builtin tools.""" + client = openai_client + tools = [tool_config] # type: ignore + + if should_raise: + with pytest.raises(ValueError): + OpenAIAgent( + name="test", + description="desc", + client=client, + model="gpt-4o", + instructions="inst", + tools=tools, # type: ignore + ) + else: + agent = OpenAIAgent( + name="test", + description="desc", + client=client, + model="gpt-4o", + instructions="inst", + tools=tools, # type: ignore + ) + assert any(t["type"] == tool_config["type"] for t in agent.tools) + + +@pytest.mark.asyncio +async def test_builtin_tool_validation_with_custom_and_builtin(openai_client: AsyncOpenAI) -> None: + """Test validation with mixed string and dictionary tools.""" + client = openai_client + tools = ["web_search_preview", {"type": "image_generation"}] # type: ignore + agent = OpenAIAgent( + name="test", + description="desc", + client=client, + model="gpt-4o", + instructions="inst", + tools=tools, # type: ignore + ) + assert any(t["type"] == "web_search_preview" for t in agent.tools) + assert any(t["type"] == "image_generation" for t in agent.tools) + + +@pytest.mark.asyncio +@skip_if_no_real_openai_key +async def test_integration_with_openai_api() -> None: + """Test basic integration with OpenAI API.""" + api_key = os.getenv("OPENAI_API_KEY") + client = AsyncOpenAI(api_key=api_key) + tools = ["web_search_preview"] # type: ignore + agent = OpenAIAgent( + name="integration", + description="desc", + client=client, + model="gpt-4o", + instructions="You are a helpful assistant.", + tools=tools, # type: ignore + ) + cancellation_token = CancellationToken() + response = await agent.on_messages( + [TextMessage(source="user", content="What is the capital of France?")], + cancellation_token, + ) + assert hasattr(response, "chat_message") + assert hasattr(response.chat_message, "content") + content = getattr(response.chat_message, "content", "") + assert content + + +@pytest.mark.asyncio +@skip_if_no_real_openai_key +async def test_integration_web_search_preview_tool() -> None: + """Test web_search_preview tool with actual API call.""" + api_key = os.getenv("OPENAI_API_KEY") + client = AsyncOpenAI(api_key=api_key) + tools = ["web_search_preview"] # type: ignore + agent = OpenAIAgent( + name="web_search_test", + description="Test agent with web search capability", + client=client, + model="gpt-4o", + instructions="You are a helpful assistant with web search capabilities. Use web search when needed.", + tools=tools, # type: ignore + ) + cancellation_token = CancellationToken() + + # Test web search functionality + response = await agent.on_messages( + [TextMessage(source="user", content="What are the latest developments in AI technology?")], + cancellation_token, + ) + assert hasattr(response, "chat_message") + assert hasattr(response.chat_message, "content") + content = getattr(response.chat_message, "content", "") + assert len(content) > 0 + + +@pytest.mark.asyncio +@skip_if_no_real_openai_key +async def test_integration_image_generation_tool() -> None: + """Test image_generation tool with actual API call.""" + api_key = os.getenv("OPENAI_API_KEY") + client = AsyncOpenAI(api_key=api_key) + tools = ["image_generation"] # type: ignore + agent = OpenAIAgent( + name="image_gen_test", + description="Test agent with image generation capability", + client=client, + model="gpt-4o", + instructions="You are a helpful assistant with image generation capabilities. Generate images when requested.", + tools=tools, # type: ignore + ) + cancellation_token = CancellationToken() + + # Test image generation functionality + response = await agent.on_messages( + [TextMessage(source="user", content="Generate an image of a beautiful sunset over mountains")], + cancellation_token, + ) + assert hasattr(response, "chat_message") + assert hasattr(response.chat_message, "content") + content = getattr(response.chat_message, "content", "") + assert len(content) > 0 + + +@pytest.mark.asyncio +@skip_if_no_real_openai_key +async def test_integration_configured_web_search_tool() -> None: + """Test web_search_preview tool with configuration using actual API call.""" + api_key = os.getenv("OPENAI_API_KEY") + client = AsyncOpenAI(api_key=api_key) + tools = [{"type": "web_search_preview", "user_location": "US", "search_context_size": 5}] # type: ignore + agent = OpenAIAgent( + name="configured_web_search_test", + description="Test agent with configured web search capability", + client=client, + model="gpt-4o", + instructions="You are a helpful assistant with configured web search capabilities.", + tools=tools, # type: ignore + ) + cancellation_token = CancellationToken() + + # Test configured web search functionality + response = await agent.on_messages( + [TextMessage(source="user", content="What's the weather like in San Francisco today?")], + cancellation_token, + ) + assert hasattr(response, "chat_message") + assert hasattr(response.chat_message, "content") + content = getattr(response.chat_message, "content", "") + assert len(content) > 0 + + +@pytest.mark.asyncio +@skip_if_no_real_openai_key +async def test_integration_configured_image_generation_tool() -> None: + """Test image_generation tool with configuration using actual API call.""" + api_key = os.getenv("OPENAI_API_KEY") + client = AsyncOpenAI(api_key=api_key) + tools = [{"type": "image_generation", "background": "white"}] # type: ignore + agent = OpenAIAgent( + name="configured_image_gen_test", + description="Test agent with configured image generation capability", + client=client, + model="gpt-4o", + instructions="You are a helpful assistant with configured image generation capabilities.", + tools=tools, # type: ignore + ) + cancellation_token = CancellationToken() + + # Test configured image generation functionality + response = await agent.on_messages( + [TextMessage(source="user", content="Create an image of a cat sitting on a white background")], + cancellation_token, + ) + assert hasattr(response, "chat_message") + assert hasattr(response.chat_message, "content") + content = getattr(response.chat_message, "content", "") + assert len(content) > 0 + + +@pytest.mark.asyncio +@skip_if_no_real_openai_key +async def test_integration_multiple_builtin_tools() -> None: + """Test multiple builtin tools together with actual API call.""" + api_key = os.getenv("OPENAI_API_KEY") + client = AsyncOpenAI(api_key=api_key) + tools = ["web_search_preview", "image_generation"] # type: ignore + agent = OpenAIAgent( + name="multi_tool_test", + description="Test agent with multiple builtin tools", + client=client, + model="gpt-4o", + instructions="You are a helpful assistant with web search and image generation capabilities.", + tools=tools, # type: ignore + ) + cancellation_token = CancellationToken() + + # Test multiple tools functionality + response = await agent.on_messages( + [ + TextMessage( + source="user", + content="Search for information about space exploration and generate an image of a rocket", + ) + ], + cancellation_token, + ) + assert hasattr(response, "chat_message") + assert hasattr(response.chat_message, "content") + content = getattr(response.chat_message, "content", "") + assert len(content) > 0 + + +@pytest.mark.asyncio +@skip_if_no_real_openai_key +async def test_integration_file_search_tool_with_vector_store() -> None: + """Test file_search tool with vector store configuration (requires actual vector store).""" + api_key = os.getenv("OPENAI_API_KEY") + + # Skip this test if no vector store ID is provided + vector_store_id = os.getenv("OPENAI_VECTOR_STORE_ID") + if not vector_store_id: + pytest.skip("OPENAI_VECTOR_STORE_ID not set; skipping file_search integration test.") + + client = AsyncOpenAI(api_key=api_key) + tools = [{"type": "file_search", "vector_store_ids": [vector_store_id]}] # type: ignore + agent = OpenAIAgent( + name="file_search_test", + description="Test agent with file search capability", + client=client, + model="gpt-4o", + instructions="You are a helpful assistant with file search capabilities.", + tools=tools, # type: ignore + ) + cancellation_token = CancellationToken() + + # Test file search functionality + response = await agent.on_messages( + [TextMessage(source="user", content="Search for documents about machine learning")], + cancellation_token, + ) + assert hasattr(response, "chat_message") + assert hasattr(response.chat_message, "content") + content = getattr(response.chat_message, "content", "") + assert len(content) > 0 + + +@pytest.mark.asyncio +@skip_if_no_real_openai_key +async def test_integration_code_interpreter_tool() -> None: + """Test code_interpreter tool with actual API call.""" + api_key = os.getenv("OPENAI_API_KEY") + client = AsyncOpenAI(api_key=api_key) + tools = [{"type": "code_interpreter", "container": "python-3.11"}] # type: ignore + agent = OpenAIAgent( + name="code_interpreter_test", + description="Test agent with code interpreter capability", + client=client, + model="gpt-4o", + instructions="You are a helpful assistant with code execution capabilities.", + tools=tools, # type: ignore + ) + cancellation_token = CancellationToken() + + # Test code interpreter functionality + response = await agent.on_messages( + [TextMessage(source="user", content="Calculate the sum of numbers from 1 to 100")], + cancellation_token, + ) + assert hasattr(response, "chat_message") + assert hasattr(response.chat_message, "content") + content = getattr(response.chat_message, "content", "") + assert len(content) > 0 + + +@pytest.mark.asyncio +@skip_if_no_real_openai_key +async def test_integration_streaming_with_builtin_tools() -> None: + """Test streaming responses with builtin tools.""" + api_key = os.getenv("OPENAI_API_KEY") + client = AsyncOpenAI(api_key=api_key) + tools = ["web_search_preview"] # type: ignore + agent = OpenAIAgent( + name="streaming_test", + description="Test agent with streaming and builtin tools", + client=client, + model="gpt-4o", + instructions="You are a helpful assistant with web search capabilities.", + tools=tools, # type: ignore + ) + cancellation_token = CancellationToken() + + # Test streaming with builtin tools + messages: list[Any] = [] + async for message in agent.on_messages_stream( + [TextMessage(source="user", content="What are the latest news about renewable energy?")], + cancellation_token, + ): + messages.append(message) + + # Verify we received some messages + assert len(messages) > 0 + # Verify at least one message has content + content_messages = [ + msg + for msg in messages + if hasattr(msg, "chat_message") + and hasattr(msg.chat_message, "content") + and getattr(msg.chat_message, "content", False) + ] + assert len(content_messages) > 0 + + +# JSON Config Tests for Built-in Tools + + +@pytest.mark.asyncio +async def test_to_config_with_string_builtin_tools() -> None: + """Test _to_config with string-based builtin tools.""" + client = AsyncOpenAI() + tools = ["web_search_preview", "image_generation"] # type: ignore + agent = OpenAIAgent( + name="config_test", + description="Test agent for config serialization", + client=client, + model="gpt-4o", + instructions="Test instructions", + tools=tools, # type: ignore + ) + + config = agent.to_config() + assert config.name == "config_test" + assert config.description == "Test agent for config serialization" + assert config.model == "gpt-4o" + assert config.instructions == "Test instructions" + assert config.tools is not None + assert len(config.tools) == 2 + + # Verify tools are serialized correctly + tool_types: list[str] = [] + for tool in config.tools: + if isinstance(tool, str): + tool_types.append(tool) + elif isinstance(tool, dict): + tool_types.append(cast(Dict[str, Any], tool)["type"]) + else: + # Handle ComponentModel case + tool_types.append(str(tool)) + assert "web_search_preview" in tool_types + assert "image_generation" in tool_types + + +@pytest.mark.asyncio +async def test_to_config_with_configured_builtin_tools() -> None: + """Test _to_config with configured builtin tools.""" + client = AsyncOpenAI() + tools = [ + {"type": "file_search", "vector_store_ids": ["vs1", "vs2"], "max_num_results": 10}, # type: ignore + {"type": "web_search_preview", "user_location": "US", "search_context_size": 5}, # type: ignore + {"type": "image_generation", "background": "white"}, # type: ignore + ] + agent = OpenAIAgent( + name="configured_test", + description="Test agent with configured tools", + client=client, + model="gpt-4o", + instructions="Test instructions", + tools=tools, # type: ignore + ) + + config = agent.to_config() + assert config.name == "configured_test" + assert config.tools is not None + assert len(config.tools) == 3 + + # Verify configured tools are serialized correctly + tool_configs = [cast(Dict[str, Any], tool) for tool in config.tools if isinstance(tool, dict)] + assert len(tool_configs) == 3 + + # Check file_search config + file_search_config = next(tool for tool in tool_configs if tool["type"] == "file_search") + assert file_search_config["vector_store_ids"] == ["vs1", "vs2"] + assert file_search_config["max_num_results"] == 10 + + # Check web_search_preview config + web_search_config = next(tool for tool in tool_configs if tool["type"] == "web_search_preview") + assert web_search_config["user_location"] == "US" + assert web_search_config["search_context_size"] == 5 + + # Check image_generation config + image_gen_config = next(tool for tool in tool_configs if tool["type"] == "image_generation") + assert image_gen_config["background"] == "white" + + +@pytest.mark.asyncio +async def test_from_config_with_string_builtin_tools() -> None: + """Test _from_config with string-based builtin tools.""" + from autogen_ext.agents.openai._openai_agent import OpenAIAgentConfig # type: ignore + + config = OpenAIAgentConfig( + name="from_config_test", + description="Test agent from config", + model="gpt-4o", + instructions="Test instructions", + tools=["web_search_preview", "image_generation"], # type: ignore + ) + agent = OpenAIAgent.from_config(config) + assert agent.name == "from_config_test" + assert agent.description == "Test agent from config" + assert agent.model == "gpt-4o" + # Verify instructions via configuration + assert agent.to_config().instructions == "Test instructions" + # Verify tools are loaded correctly + assert len(agent.tools) == 2 + tool_types = [tool["type"] for tool in agent.tools] + assert "web_search_preview" in tool_types + assert "image_generation" in tool_types + + +@pytest.mark.asyncio +async def test_from_config_with_configured_builtin_tools() -> None: + """Test _from_config with configured builtin tools.""" + from autogen_ext.agents.openai._openai_agent import OpenAIAgentConfig # type: ignore + + config = OpenAIAgentConfig( + name="configured_from_config_test", + description="Test agent with configured tools from config", + model="gpt-4o", + instructions="Test instructions", + tools=[ + {"type": "file_search", "vector_store_ids": ["vs1"]}, # type: ignore + {"type": "web_search_preview", "user_location": "US"}, # type: ignore + {"type": "image_generation", "background": "black"}, # type: ignore + ], + ) + agent = OpenAIAgent.from_config(config) + assert agent.name == "configured_from_config_test" + assert agent.model == "gpt-4o" + # Verify configured tools are loaded correctly + assert len(agent.tools) == 3 + # Check file_search + file_search_tool = next(tool for tool in agent.tools if tool["type"] == "file_search") + assert file_search_tool["vector_store_ids"] == ["vs1"] + # Check web_search_preview + web_search_tool = next(tool for tool in agent.tools if tool["type"] == "web_search_preview") + assert web_search_tool["user_location"] == "US" + # Check image_generation + image_gen_tool = next(tool for tool in agent.tools if tool["type"] == "image_generation") + assert image_gen_tool["background"] == "black" + + +@pytest.mark.asyncio +async def test_round_trip_config_serialization() -> None: + """Test round-trip serialization: agent -> config -> agent.""" + client = AsyncOpenAI() + original_tools = [ + "web_search_preview", + {"type": "file_search", "vector_store_ids": ["vs1"]}, # type: ignore + {"type": "image_generation", "background": "white"}, # type: ignore + ] + + original_agent = OpenAIAgent( + name="round_trip_test", + description="Test round-trip serialization", + client=client, + model="gpt-4o", + instructions="Test instructions", + tools=original_tools, # type: ignore + ) + + # Serialize to config + config = original_agent.to_config() + + # Deserialize back to agent + restored_agent = OpenAIAgent.from_config(config) + + # Verify basic properties + assert restored_agent.name == original_agent.name + assert restored_agent.description == original_agent.description + assert restored_agent.model == original_agent.model + orig_config = original_agent.to_config() + restored_config = restored_agent.to_config() + assert restored_config.instructions == orig_config.instructions + + # Verify tools are preserved + assert len(restored_agent.tools) == len(original_agent.tools) + + # Check that string tools are preserved + assert any(tool["type"] == "web_search_preview" for tool in restored_agent.tools) + + # Check that configured tools are preserved + file_search_tool = next(tool for tool in restored_agent.tools if tool["type"] == "file_search") + assert file_search_tool["vector_store_ids"] == ["vs1"] + + image_gen_tool = next(tool for tool in restored_agent.tools if tool["type"] == "image_generation") + assert image_gen_tool["background"] == "white" + + +@pytest.mark.asyncio +async def test_config_serialization_with_mixed_tools() -> None: + """Test config serialization with mixed string and configured tools.""" + client = AsyncOpenAI() + tools = [ + "web_search_preview", # string tool + {"type": "file_search", "vector_store_ids": ["vs1"]}, # type: ignore + "image_generation", # string tool + {"type": "code_interpreter", "container": "python-3.11"}, # type: ignore + ] + + agent = OpenAIAgent( + name="mixed_tools_test", + description="Test agent with mixed tool types", + client=client, + model="gpt-4o", + instructions="Test instructions", + tools=tools, # type: ignore + ) + + config = agent.to_config() + assert config.tools is not None + assert len(config.tools) == 4 + + # Verify all tools are serialized as dicts with "type" key + dict_tools = [cast(Dict[str, Any], tool) for tool in config.tools if isinstance(tool, dict)] + assert len(dict_tools) == 4 + + # Check that string tools are converted to dicts with "type" key + tool_types = [tool["type"] for tool in dict_tools] + assert "web_search_preview" in tool_types + assert "file_search" in tool_types + assert "image_generation" in tool_types + assert "code_interpreter" in tool_types + + # Verify configured tools preserve their configuration + file_search_config = next(tool for tool in dict_tools if tool["type"] == "file_search") + assert file_search_config["vector_store_ids"] == ["vs1"] + + code_interpreter_config = next(tool for tool in dict_tools if tool["type"] == "code_interpreter") + assert code_interpreter_config["container"] == "python-3.11" + + +@pytest.mark.asyncio +async def test_config_serialization_with_local_shell() -> None: + """Test config serialization with local_shell tool (model-restricted).""" + client = AsyncOpenAI() + tools = ["local_shell"] # type: ignore + + agent = OpenAIAgent( + name="local_shell_test", + description="Test agent with local_shell", + client=client, + model="codex-mini-latest", # Required for local_shell + instructions="Test instructions", + tools=tools, # type: ignore + ) + + config = agent.to_config() + assert config.model == "codex-mini-latest" + assert config.tools is not None + assert len(config.tools) == 1 + # Built-in tools are serialized as dicts with "type" key + assert config.tools[0] == {"type": "local_shell"} + + # Test round-trip + restored_agent = OpenAIAgent.from_config(config) + assert restored_agent.model == "codex-mini-latest" + assert len(restored_agent.tools) == 1 + assert restored_agent.tools[0]["type"] == "local_shell" + + +@pytest.mark.asyncio +async def test_config_serialization_with_complex_web_search() -> None: + """Test config serialization with complex web_search_preview configuration.""" + client = AsyncOpenAI() + tools = [ + { + "type": "web_search_preview", + "user_location": {"type": "approximate", "country": "US", "region": "CA", "city": "San Francisco"}, + "search_context_size": 10, + } + ] # type: ignore + agent = OpenAIAgent( + name="complex_web_search_test", + description="Test agent with complex web search config", + client=client, + model="gpt-4o", + instructions="Test instructions", + tools=tools, # type: ignore + ) + config = agent.to_config() + assert config.tools is not None + assert len(config.tools) == 1 + web_search_config = cast(Dict[str, Any], config.tools[0]) + assert isinstance(web_search_config, dict) + assert web_search_config["type"] == "web_search_preview" + user_location = web_search_config["user_location"] + if isinstance(user_location, dict): + assert user_location["type"] == "approximate" + assert user_location["country"] == "US" + assert user_location["region"] == "CA" + assert user_location["city"] == "San Francisco" + else: + # If user_location is a string, just check value + assert user_location == "US" + assert web_search_config["search_context_size"] == 10 + # Test round-trip + restored_agent = OpenAIAgent.from_config(config) + restored_tool = cast(Dict[str, Any], restored_agent.tools[0]) + assert restored_tool["type"] == "web_search_preview" + restored_user_location = restored_tool["user_location"] + if isinstance(restored_user_location, dict): + assert restored_user_location["type"] == "approximate" + assert restored_user_location["country"] == "US" + assert restored_user_location["region"] == "CA" + assert restored_user_location["city"] == "San Francisco" + else: + assert restored_user_location == "US" + assert restored_tool["search_context_size"] == 10 diff --git a/python/packages/autogen-ext/tests/code_executors/test_aca_dynamic_sessions.py b/python/packages/autogen-ext/tests/code_executors/test_aca_dynamic_sessions.py index 7a69a18791cc..aa13f1549f8b 100644 --- a/python/packages/autogen-ext/tests/code_executors/test_aca_dynamic_sessions.py +++ b/python/packages/autogen-ext/tests/code_executors/test_aca_dynamic_sessions.py @@ -22,6 +22,23 @@ POOL_ENDPOINT = os.getenv(ENVIRON_KEY_AZURE_POOL_ENDPOINT) +def test_session_id_preserved_if_passed() -> None: + executor = ACADynamicSessionsCodeExecutor( + pool_management_endpoint="fake-endpoint", credential=DefaultAzureCredential() + ) + session_id = "test_session_id" + executor._session_id = session_id # type: ignore[reportPrivateUsage] + assert executor._session_id == session_id # type: ignore[reportPrivateUsage] + + +def test_session_id_generated_if_not_passed() -> None: + executor = ACADynamicSessionsCodeExecutor( + pool_management_endpoint="fake-endpoint", credential=DefaultAzureCredential() + ) + assert executor._session_id is not None # type: ignore[reportPrivateUsage] + assert len(executor._session_id) > 0 # type: ignore[reportPrivateUsage] + + @pytest.mark.skipif( not POOL_ENDPOINT, reason="do not run if pool endpoint is not defined", @@ -33,6 +50,7 @@ async def test_execute_code() -> None: executor = ACADynamicSessionsCodeExecutor( pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential() ) + await executor.start() # Test single code block. code_blocks = [CodeBlock(code="import sys; print('hello world!')", language="python")] @@ -59,6 +77,56 @@ async def test_execute_code() -> None: code_blocks = [CodeBlock(code="\n".join(file_lines), language="python")] code_result = await executor.execute_code_blocks(code_blocks, cancellation_token) assert code_result.exit_code == 0 and "hello world!" in code_result.output and "200" in code_result.output + await executor.stop() + + +@pytest.mark.skipif( + not POOL_ENDPOINT, + reason="do not run if pool endpoint is not defined", +) +@pytest.mark.asyncio +async def test_execute_code_create_image() -> None: + assert POOL_ENDPOINT is not None + cancellation_token = CancellationToken() + executor = ACADynamicSessionsCodeExecutor( + pool_management_endpoint=POOL_ENDPOINT, + credential=DefaultAzureCredential(), + suppress_result_output=True, + ) + + # Test code block that creates an image. + # This code cuases the session call to return a result with the base64 encoded output + # By default, this is appended to the output + # This test verifies that suppress_result_output prevents this from happening + code_blocks = [ + CodeBlock( + code=""" +import matplotlib.pyplot as plt +import matplotlib.patches as patches + +# Create a figure and axis +fig, ax = plt.subplots(figsize=(6, 6)) + +# Add a circle +circle = patches.Circle((0.5, 0.5), 0.3, color='blue', fill=True) +ax.add_patch(circle) + + +# Set the axis limits and aspect ratio +ax.set_xlim(0, 1) +ax.set_ylim(0, 1) +ax.set_aspect('equal') +ax.axis('off') # Turn off the axis + +# Save the image to a file +plt.savefig("circle.png", bbox_inches='tight') +print("Saved to circle.png") +""", + language="python", + ), + ] + code_result = await executor.execute_code_blocks(code_blocks, cancellation_token) + assert code_result.exit_code == 0 and "base64_data" not in code_result.output @pytest.mark.skipif( @@ -72,9 +140,11 @@ async def test_azure_container_code_executor_timeout() -> None: executor = ACADynamicSessionsCodeExecutor( pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), timeout=1 ) + await executor.start() code_blocks = [CodeBlock(code="import time; time.sleep(10); print('hello world!')", language="python")] with pytest.raises(asyncio.TimeoutError): await executor.execute_code_blocks(code_blocks, cancellation_token) + await executor.stop() @pytest.mark.skipif( @@ -88,6 +158,7 @@ async def test_azure_container_code_executor_cancellation() -> None: executor = ACADynamicSessionsCodeExecutor( pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential() ) + await executor.start() code_blocks = [CodeBlock(code="import time; time.sleep(10); print('hello world!')", language="python")] coro = executor.execute_code_blocks(code_blocks, cancellation_token) @@ -97,6 +168,7 @@ async def test_azure_container_code_executor_cancellation() -> None: with pytest.raises(asyncio.CancelledError): await coro + await executor.stop() @pytest.mark.skipif( @@ -116,6 +188,7 @@ async def test_upload_files() -> None: executor = ACADynamicSessionsCodeExecutor( pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), work_dir=temp_dir ) + await executor.start() async with await open_file(os.path.join(temp_dir, test_file_1), "w") as f: await f.write(test_file_1_contents) @@ -144,6 +217,8 @@ async def test_upload_files() -> None: assert test_file_1_contents in code_result.output assert test_file_2_contents in code_result.output + await executor.stop() + @pytest.mark.skipif( not POOL_ENDPOINT, @@ -162,6 +237,7 @@ async def test_download_files() -> None: executor = ACADynamicSessionsCodeExecutor( pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), work_dir=temp_dir ) + await executor.start() code_blocks = [ CodeBlock( @@ -191,3 +267,5 @@ async def test_download_files() -> None: async with await open_file(os.path.join(temp_dir, test_file_2), "r") as f: content = await f.read() assert test_file_2_contents in content + + await executor.stop() diff --git a/python/packages/autogen-ext/tests/code_executors/test_aca_user_defined_functions.py b/python/packages/autogen-ext/tests/code_executors/test_aca_user_defined_functions.py index 27adbee72cdf..b0a50837ece4 100644 --- a/python/packages/autogen-ext/tests/code_executors/test_aca_user_defined_functions.py +++ b/python/packages/autogen-ext/tests/code_executors/test_aca_user_defined_functions.py @@ -61,6 +61,7 @@ async def test_azure_can_load_function_with_reqs() -> None: azure_executor = ACADynamicSessionsCodeExecutor( pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), functions=[load_data] ) + await azure_executor.start() # ACADynamicSessionsCodeExecutor doesn't use the functions module import code = """import polars @@ -77,6 +78,8 @@ async def test_azure_can_load_function_with_reqs() -> None: assert azure_result.output == "John\n" assert azure_result.exit_code == 0 + await azure_executor.stop() + @pytest.mark.skipif( not POOL_ENDPOINT, @@ -90,6 +93,8 @@ async def test_azure_can_load_function() -> None: azure_executor = ACADynamicSessionsCodeExecutor( pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), functions=[add_two_numbers] ) + await azure_executor.start() + # ACADynamicSessionsCodeExecutor doesn't use the functions module import code = """print(add_two_numbers(1, 2))""" @@ -102,6 +107,8 @@ async def test_azure_can_load_function() -> None: assert azure_result.output == "3\n" assert azure_result.exit_code == 0 + await azure_executor.stop() + @pytest.mark.skipif( not POOL_ENDPOINT, @@ -116,6 +123,8 @@ async def test_azure_fails_for_function_incorrect_import() -> None: credential=DefaultAzureCredential(), functions=[function_incorrect_import], ) + await azure_executor.start() + code = """function_incorrect_import()""" with pytest.raises(ValueError): @@ -126,6 +135,8 @@ async def test_azure_fails_for_function_incorrect_import() -> None: cancellation_token=cancellation_token, ) + await azure_executor.stop() + @pytest.mark.skipif( not POOL_ENDPOINT, @@ -138,6 +149,7 @@ async def test_azure_fails_for_function_incorrect_dep() -> None: azure_executor = ACADynamicSessionsCodeExecutor( pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), functions=[function_incorrect_dep] ) + await azure_executor.start() code = """function_incorrect_dep()""" with pytest.raises(ValueError): @@ -148,6 +160,8 @@ async def test_azure_fails_for_function_incorrect_dep() -> None: cancellation_token=cancellation_token, ) + await azure_executor.stop() + def test_azure_formatted_prompt() -> None: assert_str = '''def add_two_numbers(a: int, b: int) -> int: @@ -200,6 +214,8 @@ def add_two_numbers(a: int, b: int) -> int: azure_executor = ACADynamicSessionsCodeExecutor( pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), functions=[func] ) + await azure_executor.start() + code = """print(add_two_numbers(1, 2))""" azure_result = await azure_executor.execute_code_blocks( @@ -211,6 +227,8 @@ def add_two_numbers(a: int, b: int) -> int: assert azure_result.output == "3\n" assert azure_result.exit_code == 0 + await azure_executor.stop() + @pytest.mark.skipif( not POOL_ENDPOINT, @@ -231,6 +249,8 @@ def add_two_numbers(a: int, b: int) -> int: azure_executor = ACADynamicSessionsCodeExecutor( pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), functions=[func] ) + await azure_executor.start() + code = """print(add_two_numbers(object(), False))""" azure_result = await azure_executor.execute_code_blocks( @@ -242,3 +262,5 @@ def add_two_numbers(a: int, b: int) -> int: # result.output = result.output.encode().decode('unicode_escape') assert "TypeError: unsupported operand type(s) for +:" in azure_result.output assert azure_result.exit_code == 1 + + await azure_executor.stop() diff --git a/python/packages/autogen-ext/tests/code_executors/test_commandline_code_executor.py b/python/packages/autogen-ext/tests/code_executors/test_commandline_code_executor.py index 2a5cedcb5225..6ba30da76f15 100644 --- a/python/packages/autogen-ext/tests/code_executors/test_commandline_code_executor.py +++ b/python/packages/autogen-ext/tests/code_executors/test_commandline_code_executor.py @@ -5,11 +5,14 @@ import os import platform import shutil +import subprocess import sys import tempfile +import types import venv from pathlib import Path from typing import AsyncGenerator, TypeAlias +from unittest.mock import patch import pytest import pytest_asyncio @@ -21,6 +24,85 @@ HAS_POWERSHELL: bool = platform.system() == "Windows" and ( shutil.which("powershell") is not None or shutil.which("pwsh") is not None ) +IS_MACOS: bool = platform.system() == "Darwin" +IS_UV_VENV: bool = ( + lambda: ( + ( + lambda venv_path: ( + False + if not venv_path + else ( + False + if not os.path.isfile(os.path.join(venv_path, "pyvenv.cfg")) + else ( + subprocess.run( + ["grep", "-q", "^uv = ", os.path.join(venv_path, "pyvenv.cfg")], + check=False, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ).returncode + == 0 + ) + ) + ) + )(os.environ.get("VIRTUAL_ENV")) + ) +)() +HAS_UV: bool = shutil.which("uv") is not None + + +def create_venv_with_uv(env_dir: str) -> types.SimpleNamespace: + try: + subprocess.run( + ["uv", "venv", env_dir], + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + except subprocess.CalledProcessError as e: + error_message = f"uv virtual env creation failed with error code {e.returncode}:\n" + error_message += f" cmd:\n{e.stdout.decode()}\n" + error_message += f" stderr:\n{e.stderr}\n" + error_message += f" stdout:\n{e.stdout}" + raise RuntimeError(error_message) from e + except Exception as e: + raise RuntimeError(f"Failed to create uv virtual env: {e}") from e + + # create a venv.EnvBuilder context + if platform.system() == "Windows": + bin_name = "Scripts" + exe_suffix = ".exe" + else: + bin_name = "bin" + exe_suffix = "" + + bin_path = os.path.join(env_dir, bin_name) + python_executable = os.path.join(bin_path, f"python{exe_suffix}") + py_version_short = f"{sys.version_info.major}.{sys.version_info.minor}" + lib_path = os.path.join(env_dir, "lib", f"python{py_version_short}", "site-packages") + if not os.path.exists(lib_path): + lib_path_fallback = os.path.join(env_dir, "lib") + if os.path.exists(lib_path_fallback): + lib_path = lib_path_fallback + else: + raise RuntimeError(f"Failed to find site-packages in {lib_path} or {lib_path_fallback}") + + context = types.SimpleNamespace( + env_dir=env_dir, + env_name=os.path.basename(env_dir), + prompt=f"({os.path.basename(env_dir)}) ", + executable=python_executable, + python_dir=os.path.dirname(python_executable), + python_exe=os.path.basename(python_executable), + inc_path=os.path.join(env_dir, "include"), + lib_path=lib_path, # site-packages + bin_path=bin_path, # bin or Scripts + bin_name=bin_name, # bin or Scripts + env_exe=python_executable, + env_exec_cmd=python_executable, + ) + + return context @pytest_asyncio.fixture(scope="function") # type: ignore @@ -28,7 +110,9 @@ async def executor_and_temp_dir( request: pytest.FixtureRequest, ) -> AsyncGenerator[tuple[LocalCommandLineCodeExecutor, str], None]: with tempfile.TemporaryDirectory() as temp_dir: - yield LocalCommandLineCodeExecutor(work_dir=temp_dir), temp_dir + executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, cleanup_temp_files=False) + await executor.start() + yield executor, temp_dir ExecutorFixture: TypeAlias = tuple[LocalCommandLineCodeExecutor, str] @@ -98,6 +182,7 @@ async def test_commandline_code_executor_cancellation() -> None: with tempfile.TemporaryDirectory() as temp_dir: cancellation_token = CancellationToken() executor = LocalCommandLineCodeExecutor(work_dir=temp_dir) + await executor.start() # Write code that sleep for 10 seconds and then write "hello world!" # to a file. code = """import time @@ -166,6 +251,10 @@ async def test_valid_relative_path(executor_and_temp_dir: ExecutorFixture) -> No @pytest.mark.asyncio +@pytest.mark.skipif( + IS_MACOS and IS_UV_VENV, + reason="uv-venv is not supported on macOS.", +) async def test_local_executor_with_custom_venv() -> None: with tempfile.TemporaryDirectory() as temp_dir: env_builder = venv.EnvBuilder(with_pip=True) @@ -173,6 +262,8 @@ async def test_local_executor_with_custom_venv() -> None: env_builder_context = env_builder.ensure_directories(temp_dir) executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, virtual_env_context=env_builder_context) + await executor.start() + code_blocks = [ # https://stackoverflow.com/questions/1871549/how-to-determine-if-python-is-running-inside-a-virtualenv CodeBlock(code="import sys; print(sys.prefix != sys.base_prefix)", language="python"), @@ -185,6 +276,10 @@ async def test_local_executor_with_custom_venv() -> None: @pytest.mark.asyncio +@pytest.mark.skipif( + IS_MACOS and IS_UV_VENV, + reason="uv-venv is not supported on macOS.", +) async def test_local_executor_with_custom_venv_in_local_relative_path() -> None: relative_folder_path = "tmp_dir" try: @@ -197,6 +292,8 @@ async def test_local_executor_with_custom_venv_in_local_relative_path() -> None: env_builder_context = env_builder.ensure_directories(env_path) executor = LocalCommandLineCodeExecutor(work_dir=relative_folder_path, virtual_env_context=env_builder_context) + await executor.start() + code_blocks = [ CodeBlock(code="import sys; print(sys.executable)", language="python"), ] @@ -213,13 +310,75 @@ async def test_local_executor_with_custom_venv_in_local_relative_path() -> None: shutil.rmtree(relative_folder_path) -def test_serialize_deserialize() -> None: +@pytest.mark.asyncio +@pytest.mark.skipif( + not HAS_UV, + reason="uv is not installed.", +) +async def test_local_executor_with_custom_uv_venv() -> None: + with tempfile.TemporaryDirectory() as temp_dir: + env_builder_context = create_venv_with_uv(temp_dir) + + executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, virtual_env_context=env_builder_context) + await executor.start() + + code_blocks = [ + # https://stackoverflow.com/questions/1871549/how-to-determine-if-python-is-running-inside-a-virtualenv + CodeBlock(code="import sys; print(sys.prefix != sys.base_prefix)", language="python"), + ] + cancellation_token = CancellationToken() + result = await executor.execute_code_blocks(code_blocks, cancellation_token=cancellation_token) + + assert result.exit_code == 0 + assert result.output.strip() == "True" + + +@pytest.mark.asyncio +@pytest.mark.skipif( + not HAS_UV, + reason="uv is not installed.", +) +async def test_local_executor_with_custom_uv_venv_in_local_relative_path() -> None: + relative_folder_path = "tmp_dir" + try: + if not os.path.isdir(relative_folder_path): + os.mkdir(relative_folder_path) + + env_path = os.path.join(relative_folder_path, ".venv") + env_builder_context = create_venv_with_uv(env_path) + + executor = LocalCommandLineCodeExecutor(work_dir=relative_folder_path, virtual_env_context=env_builder_context) + await executor.start() + + code_blocks = [ + CodeBlock(code="import sys; print(sys.executable)", language="python"), + ] + cancellation_token = CancellationToken() + result = await executor.execute_code_blocks(code_blocks, cancellation_token=cancellation_token) + + assert result.exit_code == 0 + + # Check if the expected venv has been used + bin_path = os.path.abspath(env_builder_context.bin_path) + assert Path(result.output.strip()).parent.samefile(bin_path) + finally: + if os.path.isdir(relative_folder_path): + shutil.rmtree(relative_folder_path) + + +@pytest.mark.asyncio +async def test_serialize_deserialize() -> None: with tempfile.TemporaryDirectory() as temp_dir: executor = LocalCommandLineCodeExecutor(work_dir=temp_dir) + await executor.start() executor_config = executor.dump_component() loaded_executor = LocalCommandLineCodeExecutor.load_component(executor_config) + await loaded_executor.start() assert executor.work_dir == loaded_executor.work_dir + await executor.stop() + await loaded_executor.stop() + @pytest.mark.asyncio @pytest.mark.windows @@ -241,3 +400,48 @@ async def test_ps1_script(executor_and_temp_dir: ExecutorFixture) -> None: assert result.exit_code == 0 assert "hello from powershell!" in result.output assert result.code_file is not None + + +@pytest.mark.asyncio +async def test_cleanup_temp_files_behavior() -> None: + with tempfile.TemporaryDirectory() as temp_dir: + # Test with cleanup_temp_files=True (default) + executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, cleanup_temp_files=True) + await executor.start() + cancellation_token = CancellationToken() + code_blocks = [CodeBlock(code="print('cleanup test')", language="python")] + result = await executor.execute_code_blocks(code_blocks, cancellation_token) + assert result.exit_code == 0 + assert "cleanup test" in result.output + # The code file should have been deleted + assert result.code_file is not None + assert not Path(result.code_file).exists() + + # Test with cleanup_temp_files=False + executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, cleanup_temp_files=False) + await executor.start() + cancellation_token = CancellationToken() + code_blocks = [CodeBlock(code="print('no cleanup')", language="python")] + result = await executor.execute_code_blocks(code_blocks, cancellation_token) + assert result.exit_code == 0 + assert "no cleanup" in result.output + # The code file should still exist + assert result.code_file is not None + assert Path(result.code_file).exists() + + +@pytest.mark.asyncio +async def test_cleanup_temp_files_oserror(caplog: pytest.LogCaptureFixture) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, cleanup_temp_files=True) + await executor.start() + cancellation_token = CancellationToken() + code_blocks = [CodeBlock(code="print('cleanup test')", language="python")] + + # Patch Path.unlink to raise OSError for this test + with patch("pathlib.Path.unlink", side_effect=OSError("Mocked OSError")): + with caplog.at_level("ERROR"): + await executor.execute_code_blocks(code_blocks, cancellation_token) + # The code file should have been attempted to be deleted and failed + assert any("Failed to delete temporary file" in record.message for record in caplog.records) + assert any("Mocked OSError" in record.message for record in caplog.records) diff --git a/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py b/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py index 24ccdc218663..81c890efa643 100644 --- a/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py +++ b/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py @@ -1,6 +1,7 @@ # mypy: disable-error-code="no-any-unimported" import asyncio import os +import shutil import sys import tempfile from pathlib import Path @@ -32,7 +33,7 @@ def docker_tests_enabled() -> bool: return False -@pytest_asyncio.fixture(scope="function") # type: ignore +@pytest_asyncio.fixture(scope="module") # type: ignore async def executor_and_temp_dir( request: pytest.FixtureRequest, ) -> AsyncGenerator[tuple[DockerCommandLineCodeExecutor, str], None]: @@ -47,9 +48,20 @@ async def executor_and_temp_dir( ExecutorFixture: TypeAlias = tuple[DockerCommandLineCodeExecutor, str] +@pytest_asyncio.fixture(scope="function") # type: ignore +async def cleanup_temp_dir(executor_and_temp_dir: ExecutorFixture) -> AsyncGenerator[None, None]: + _executor, temp_dir = executor_and_temp_dir + for file in Path(temp_dir).iterdir(): + if file.is_file(): + file.unlink() + elif file.is_dir(): + shutil.rmtree(file) + yield None + + @pytest.mark.asyncio @pytest.mark.parametrize("executor_and_temp_dir", ["docker"], indirect=True) -async def test_execute_code(executor_and_temp_dir: ExecutorFixture) -> None: +async def test_execute_code(executor_and_temp_dir: ExecutorFixture, cleanup_temp_dir: None) -> None: executor, _temp_dir = executor_and_temp_dir cancellation_token = CancellationToken() @@ -97,7 +109,9 @@ async def test_execute_code(executor_and_temp_dir: ExecutorFixture) -> None: @pytest.mark.asyncio @pytest.mark.parametrize("executor_and_temp_dir", ["docker"], indirect=True) -async def test_commandline_code_executor_timeout(executor_and_temp_dir: ExecutorFixture) -> None: +async def test_commandline_code_executor_timeout( + executor_and_temp_dir: ExecutorFixture, cleanup_temp_dir: None +) -> None: _executor, temp_dir = executor_and_temp_dir cancellation_token = CancellationToken() code_blocks = [CodeBlock(code="import time; time.sleep(10); print('hello world!')", language="python")] @@ -110,7 +124,9 @@ async def test_commandline_code_executor_timeout(executor_and_temp_dir: Executor @pytest.mark.asyncio @pytest.mark.parametrize("executor_and_temp_dir", ["docker"], indirect=True) -async def test_commandline_code_executor_cancellation(executor_and_temp_dir: ExecutorFixture) -> None: +async def test_commandline_code_executor_cancellation( + executor_and_temp_dir: ExecutorFixture, cleanup_temp_dir: None +) -> None: _executor, temp_dir = executor_and_temp_dir cancellation_token = CancellationToken() # Write code that sleep for 10 seconds and then write "hello world!" @@ -137,7 +153,7 @@ async def test_commandline_code_executor_cancellation(executor_and_temp_dir: Exe @pytest.mark.asyncio @pytest.mark.parametrize("executor_and_temp_dir", ["docker"], indirect=True) -async def test_invalid_relative_path(executor_and_temp_dir: ExecutorFixture) -> None: +async def test_invalid_relative_path(executor_and_temp_dir: ExecutorFixture, cleanup_temp_dir: None) -> None: executor, _temp_dir = executor_and_temp_dir cancellation_token = CancellationToken() code = """# filename: /tmp/test.py @@ -152,7 +168,7 @@ async def test_invalid_relative_path(executor_and_temp_dir: ExecutorFixture) -> @pytest.mark.asyncio @pytest.mark.parametrize("executor_and_temp_dir", ["docker"], indirect=True) -async def test_valid_relative_path(executor_and_temp_dir: ExecutorFixture) -> None: +async def test_valid_relative_path(executor_and_temp_dir: ExecutorFixture, cleanup_temp_dir: None) -> None: executor, temp_dir_str = executor_and_temp_dir cancellation_token = CancellationToken() @@ -243,6 +259,142 @@ async def test_docker_commandline_code_executor_extra_args() -> None: async def test_docker_commandline_code_executor_serialization() -> None: with tempfile.TemporaryDirectory() as temp_dir: executor = DockerCommandLineCodeExecutor(work_dir=temp_dir) - loaded_executor = DockerCommandLineCodeExecutor.load_component(executor.dump_component()) + + executor_config = executor.dump_component() + loaded_executor = DockerCommandLineCodeExecutor.load_component(executor_config) + assert executor.bind_dir == loaded_executor.bind_dir assert executor.timeout == loaded_executor.timeout + + +def test_invalid_timeout() -> None: + with pytest.raises(ValueError, match="Timeout must be greater than or equal to 1."): + _ = DockerCommandLineCodeExecutor(timeout=0) + + +@pytest.mark.asyncio +async def test_directory_not_initialized() -> None: + executor = DockerCommandLineCodeExecutor() + with pytest.raises(RuntimeError, match="Working directory not properly initialized"): + _ = executor.work_dir + + +@pytest.mark.asyncio +@pytest.mark.parametrize("executor_and_temp_dir", ["docker"], indirect=True) +async def test_error_wrong_path(executor_and_temp_dir: ExecutorFixture, cleanup_temp_dir: None) -> None: + if not docker_tests_enabled(): + pytest.skip("Docker tests are disabled") + + executor, _ = executor_and_temp_dir + cancellation_token = CancellationToken() + code_blocks = [ + CodeBlock( + code="""with open("/nonexistent_dir/test.txt", "w") as f: + f.write("hello world!")""", + language="python", + ) + ] + result = await executor.execute_code_blocks(code_blocks, cancellation_token) + assert result.exit_code != 0 + assert "No such file or directory" in result.output + + +@pytest.mark.asyncio +async def test_deprecated_warning() -> None: + if not docker_tests_enabled(): + pytest.skip("Docker tests are disabled") + + with pytest.warns(DeprecationWarning, match="Using the current directory as work_dir is deprecated."): + async with DockerCommandLineCodeExecutor(work_dir=".") as executor: + await executor.start() + cancellation_token = CancellationToken() + code_block = CodeBlock(code='echo "hello world!"', language="sh") + result = await executor.execute_code_blocks([code_block], cancellation_token) + assert result.exit_code == 0 + assert "hello world!" in result.output + + +@pytest.mark.asyncio +async def test_directory_creation_cleanup() -> None: + executor = DockerCommandLineCodeExecutor(timeout=60, work_dir=None) + + await executor.start() + + directory = executor.work_dir + assert directory.is_dir() + + await executor.stop() + + assert not Path(directory).exists() + + +@pytest.mark.asyncio +async def test_delete_tmp_files() -> None: + if not docker_tests_enabled(): + pytest.skip("Docker tests are disabled") + + with tempfile.TemporaryDirectory() as temp_dir: + # Test with delete_tmp_files=False (default) + async with DockerCommandLineCodeExecutor(work_dir=temp_dir) as executor: + cancellation_token = CancellationToken() + code_blocks = [CodeBlock(code="print('test output')", language="python")] + result = await executor.execute_code_blocks(code_blocks, cancellation_token) + assert result.exit_code == 0 + assert result.code_file is not None + # Verify file exists after execution + assert Path(result.code_file).exists() + + # Test with delete_tmp_files=True + async with DockerCommandLineCodeExecutor(work_dir=temp_dir, delete_tmp_files=True) as executor: + cancellation_token = CancellationToken() + code_blocks = [CodeBlock(code="print('test output')", language="python")] + result = await executor.execute_code_blocks(code_blocks, cancellation_token) + assert result.exit_code == 0 + assert result.code_file is not None + # Verify file is deleted after execution + assert not Path(result.code_file).exists() + + # Test with multiple code blocks + code_blocks = [ + CodeBlock(code="print('first block')", language="python"), + CodeBlock(code="print('second block')", language="python"), + ] + result = await executor.execute_code_blocks(code_blocks, cancellation_token) + assert result.exit_code == 0 + assert result.code_file is not None + # Verify files are deleted after execution + assert not Path(result.code_file).exists() + + # Test deletion with execution error + code_blocks = [CodeBlock(code="raise Exception('test error')", language="python")] + result = await executor.execute_code_blocks(code_blocks, cancellation_token) + assert result.exit_code != 0 + assert result.code_file is not None + # Verify file is deleted even after error + assert not Path(result.code_file).exists() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("executor_and_temp_dir", ["docker"], indirect=True) +async def test_docker_commandline_code_executor_with_multiple_tasks( + executor_and_temp_dir: ExecutorFixture, cleanup_temp_dir: None +) -> None: + if not docker_tests_enabled(): + pytest.skip("Docker tests are disabled") + + async def run_cancellation_scenario(executor: DockerCommandLineCodeExecutor) -> None: + token = CancellationToken() + code_block = CodeBlock(language="bash", code="sleep 10") + exec_task = asyncio.create_task(executor.execute_code_blocks([code_block], cancellation_token=token)) + await asyncio.sleep(1) + token.cancel() + try: + await exec_task + except asyncio.CancelledError: + pass + + def run_scenario_in_new_loop(executor_instance: DockerCommandLineCodeExecutor) -> None: + asyncio.run(run_cancellation_scenario(executor_instance)) + + executor, _ = executor_and_temp_dir + await asyncio.get_running_loop().run_in_executor(None, run_scenario_in_new_loop, executor) diff --git a/python/packages/autogen-ext/tests/code_executors/test_docker_jupyter_code_executor.py b/python/packages/autogen-ext/tests/code_executors/test_docker_jupyter_code_executor.py new file mode 100644 index 000000000000..ad4460a78469 --- /dev/null +++ b/python/packages/autogen-ext/tests/code_executors/test_docker_jupyter_code_executor.py @@ -0,0 +1,174 @@ +import inspect +import os +import tempfile +from pathlib import Path +from typing import AsyncGenerator, TypeAlias + +import pytest +import pytest_asyncio +from autogen_core import CancellationToken +from autogen_core.code_executor import CodeBlock +from autogen_ext.code_executors.docker_jupyter import ( + DockerJupyterCodeExecutor, + DockerJupyterServer, +) + + +def docker_tests_enabled() -> bool: + if os.environ.get("SKIP_DOCKER", "unset").lower() == "true": + return False + + try: + import docker + from docker.errors import DockerException + except ImportError: + return False + + try: + client = docker.from_env() + client.ping() # type: ignore + return True + except DockerException: + return False + + +@pytest_asyncio.fixture(scope="function") # type: ignore +async def executor_and_temp_dir( + request: pytest.FixtureRequest, +) -> AsyncGenerator[tuple[DockerJupyterCodeExecutor, str], None]: + if not docker_tests_enabled(): + pytest.skip("Docker tests are disabled") + + with tempfile.TemporaryDirectory() as temp_dir: + async with DockerJupyterServer(bind_dir=temp_dir) as jupyter_server: + async with DockerJupyterCodeExecutor(jupyter_server=jupyter_server) as executor: + yield executor, temp_dir + + +ExecutorFixture: TypeAlias = tuple[DockerJupyterCodeExecutor, str] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("executor_and_temp_dir", ["docker"], indirect=True) +async def test_execute_code(executor_and_temp_dir: ExecutorFixture) -> None: + executor, _temp_dir = executor_and_temp_dir + # Test single code block. + code_blocks = [CodeBlock(code="import sys; print('hello world!')", language="python")] + code_result = await executor.execute_code_blocks(code_blocks, cancellation_token=CancellationToken()) + assert code_result.exit_code == 0 and "hello world!" in code_result.output + + # Test multiple code blocks. + code_blocks = [ + CodeBlock(code="import sys; print('hello world!')", language="python"), + CodeBlock(code="a = 100 + 100; print(a)", language="python"), + ] + code_result = await executor.execute_code_blocks(code_blocks, cancellation_token=CancellationToken()) + assert code_result.exit_code == 0 and "hello world!" in code_result.output and "200" in code_result.output + + # Test running code. + file_lines = ["import sys", "print('hello world!')", "a = 100 + 100", "print(a)"] + code_blocks = [CodeBlock(code="\n".join(file_lines), language="python")] + code_result = await executor.execute_code_blocks(code_blocks, cancellation_token=CancellationToken()) + assert code_result.exit_code == 0 and "hello world!" in code_result.output and "200" in code_result.output + + +@pytest.mark.asyncio +@pytest.mark.parametrize("executor_and_temp_dir", ["docker"], indirect=True) +async def test_execute_code_and_persist_variable(executor_and_temp_dir: ExecutorFixture) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + async with DockerJupyterServer(bind_dir=temp_dir) as jupyter_server: + async with DockerJupyterCodeExecutor(jupyter_server=jupyter_server) as executor: + code_blocks_first = [ + CodeBlock(code="a = 100 + 100; print(a)", language="python"), + ] + code_result_first = await executor.execute_code_blocks( + code_blocks_first, cancellation_token=CancellationToken() + ) + assert code_result_first.exit_code == 0 and "200" in code_result_first.output + code_blocks_second = [ + CodeBlock(code="b = a + 100 ; print(b)", language="python"), + ] + code_result_second = await executor.execute_code_blocks( + code_blocks_second, cancellation_token=CancellationToken() + ) + assert code_result_second.exit_code == 0 and "300" in code_result_second.output + + +@pytest.mark.asyncio +@pytest.mark.parametrize("executor_and_temp_dir", ["docker"], indirect=True) +async def test_timeout(executor_and_temp_dir: ExecutorFixture) -> None: + code_blocks = [CodeBlock(code="import time; time.sleep(10); print('hello world!')", language="python")] + async with DockerJupyterServer() as jupyter_server: + async with DockerJupyterCodeExecutor(jupyter_server=jupyter_server, timeout=1) as executor: + code_result = await executor.execute_code_blocks( + code_blocks=code_blocks, cancellation_token=CancellationToken() + ) + + assert code_result.exit_code and "Timeout" in code_result.output + + +@pytest.mark.asyncio +@pytest.mark.parametrize("executor_and_temp_dir", ["docker"], indirect=True) +async def test_canncellation(executor_and_temp_dir: ExecutorFixture) -> None: + _executor, temp_dir = executor_and_temp_dir + # Write code that sleep for 10 seconds and then write "hello world!" + # to a file. + code = """import time, os +time.sleep(10) +with open("hello.txt", "w") as f: + f.write("hello world!") + """ + code_blocks = [CodeBlock(code=code, language="python")] + code_result = await _executor.execute_code_blocks(code_blocks, cancellation_token=CancellationToken()) + # Check if the file was created + hello_file_path = Path(temp_dir) / "hello.txt" + assert hello_file_path.exists() and code_result.exit_code == 0 + + +@pytest.mark.asyncio +async def test_start_stop() -> None: + if not docker_tests_enabled(): + pytest.skip("Docker tests are disabled") + with tempfile.TemporaryDirectory() as temp_dir: + jupyter_server = DockerJupyterServer(bind_dir=temp_dir) + executor = DockerJupyterCodeExecutor(jupyter_server=jupyter_server) + await executor.start() + await executor.stop() + + +@pytest.mark.asyncio +async def test_invalid_timeout() -> None: + if not docker_tests_enabled(): + pytest.skip("Docker tests are disabled") + with pytest.raises(ValueError, match="Timeout must be greater than or equal to 1."): + with tempfile.TemporaryDirectory() as temp_dir: + async with DockerJupyterServer(bind_dir=temp_dir) as jupyter_server: + _ = DockerJupyterCodeExecutor(jupyter_server=jupyter_server, timeout=0) + + +@pytest.mark.asyncio +async def test_execute_code_with_image_output() -> None: + if not docker_tests_enabled(): + pytest.skip("Docker tests are disabled") + with tempfile.TemporaryDirectory() as temp_dir: + async with DockerJupyterServer(bind_dir=temp_dir) as jupyter_server: + async with DockerJupyterCodeExecutor(jupyter_server=jupyter_server) as executor: + code_blocks = [ + CodeBlock( + code=inspect.cleandoc(""" + !pip install pillow + from PIL import Image, ImageDraw + img = Image.new("RGB", (100, 100), color="white") + draw = ImageDraw.Draw(img) + draw.rectangle((10, 10, 90, 90), outline="black", fill="blue") + display(img) + """), + language="python", + ) + ] + + code_result = await executor.execute_code_blocks(code_blocks, cancellation_token=CancellationToken()) + assert len(code_result.output_files) == 1 + assert code_result.exit_code == 0 + assert "" in code_result.output + assert str(Path(code_result.output_files[0]).parent) == temp_dir diff --git a/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py b/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py index 2ff0f1cb25aa..b6789d0b5e41 100644 --- a/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py +++ b/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py @@ -11,14 +11,17 @@ @pytest.mark.asyncio async def test_execute_code(tmp_path: Path) -> None: async with JupyterCodeExecutor(output_dir=tmp_path) as executor: + await executor.start() code_blocks = [CodeBlock(code="import sys; print('hello world!')", language="python")] code_result = await executor.execute_code_blocks(code_blocks, CancellationToken()) assert code_result == JupyterCodeResult(exit_code=0, output="hello world!\n", output_files=[]) + await executor.stop() @pytest.mark.asyncio async def test_execute_code_error(tmp_path: Path) -> None: async with JupyterCodeExecutor(output_dir=tmp_path) as executor: + await executor.start() code_blocks = [CodeBlock(code="print(undefined_variable)", language="python")] code_result = await executor.execute_code_blocks(code_blocks, CancellationToken()) assert code_result == JupyterCodeResult( @@ -33,22 +36,26 @@ async def test_execute_code_error(tmp_path: Path) -> None: """), output_files=[], ) + await executor.stop() @pytest.mark.asyncio async def test_execute_multiple_code_blocks(tmp_path: Path) -> None: async with JupyterCodeExecutor(output_dir=tmp_path) as executor: + await executor.start() code_blocks = [ CodeBlock(code="import sys; print('hello world!')", language="python"), CodeBlock(code="a = 100 + 100; print(a)", language="python"), ] code_result = await executor.execute_code_blocks(code_blocks, CancellationToken()) assert code_result == JupyterCodeResult(exit_code=0, output="hello world!\n\n200\n", output_files=[]) + await executor.stop() @pytest.mark.asyncio async def test_depedent_executions(tmp_path: Path) -> None: async with JupyterCodeExecutor(output_dir=tmp_path) as executor: + await executor.start() code_blocks_1 = [CodeBlock(code="a = 'hello world!'", language="python")] code_blocks_2 = [ CodeBlock(code="print(a)", language="python"), @@ -56,11 +63,13 @@ async def test_depedent_executions(tmp_path: Path) -> None: await executor.execute_code_blocks(code_blocks_1, CancellationToken()) code_result = await executor.execute_code_blocks(code_blocks_2, CancellationToken()) assert code_result == JupyterCodeResult(exit_code=0, output="hello world!\n", output_files=[]) + await executor.stop() @pytest.mark.asyncio async def test_execute_multiple_code_blocks_error(tmp_path: Path) -> None: async with JupyterCodeExecutor(output_dir=tmp_path) as executor: + await executor.start() code_blocks = [ CodeBlock(code="import sys; print('hello world!')", language="python"), CodeBlock(code="a = 100 + 100; print(a); print(undefined_variable)", language="python"), @@ -82,30 +91,37 @@ async def test_execute_multiple_code_blocks_error(tmp_path: Path) -> None: """), output_files=[], ) + await executor.stop() @pytest.mark.asyncio async def test_execute_code_after_restart(tmp_path: Path) -> None: async with JupyterCodeExecutor(output_dir=tmp_path) as executor: + await executor.start() await executor.restart() code_blocks = [CodeBlock(code="import sys; print('hello world!')", language="python")] code_result = await executor.execute_code_blocks(code_blocks, CancellationToken()) assert code_result == JupyterCodeResult(exit_code=0, output="hello world!\n", output_files=[]) + await executor.stop() @pytest.mark.asyncio async def test_commandline_code_executor_timeout(tmp_path: Path) -> None: async with JupyterCodeExecutor(output_dir=tmp_path, timeout=2) as executor: + await executor.start() code_blocks = [CodeBlock(code="import time; time.sleep(10); print('hello world!')", language="python")] with pytest.raises(asyncio.TimeoutError): await executor.execute_code_blocks(code_blocks, CancellationToken()) + await executor.stop() + @pytest.mark.asyncio async def test_commandline_code_executor_cancellation(tmp_path: Path) -> None: async with JupyterCodeExecutor(output_dir=tmp_path) as executor: + await executor.start() code_blocks = [CodeBlock(code="import time; time.sleep(10); print('hello world!')", language="python")] cancellation_token = CancellationToken() @@ -117,10 +133,13 @@ async def test_commandline_code_executor_cancellation(tmp_path: Path) -> None: with pytest.raises(asyncio.CancelledError): await code_result_coroutine + await executor.stop() + @pytest.mark.asyncio async def test_execute_code_with_image_output(tmp_path: Path) -> None: async with JupyterCodeExecutor(output_dir=tmp_path) as executor: + await executor.start() code_blocks = [ CodeBlock( code=inspect.cleandoc(""" @@ -144,10 +163,13 @@ async def test_execute_code_with_image_output(tmp_path: Path) -> None: ) assert code_result.output_files[0].parent == tmp_path + await executor.stop() + @pytest.mark.asyncio async def test_execute_code_with_html_output(tmp_path: Path) -> None: async with JupyterCodeExecutor(output_dir=tmp_path) as executor: + await executor.start() code_blocks = [ CodeBlock( code=inspect.cleandoc(""" @@ -168,11 +190,38 @@ async def test_execute_code_with_html_output(tmp_path: Path) -> None: ) assert code_result.output_files[0].parent == tmp_path + await executor.stop() + @pytest.mark.asyncio async def test_jupyter_code_executor_serialization(tmp_path: Path) -> None: executor = JupyterCodeExecutor(output_dir=tmp_path) + await executor.start() serialized = executor.dump_component() loaded_executor = JupyterCodeExecutor.load_component(serialized) + await loaded_executor.start() assert isinstance(loaded_executor, JupyterCodeExecutor) + + await loaded_executor.stop() + await executor.stop() + + +def test_invalid_timeout() -> None: + with pytest.raises(ValueError, match="Timeout must be greater than or equal to 1."): + _ = JupyterCodeExecutor(timeout=0) + + +@pytest.mark.asyncio +async def test_deprecation_output_dir() -> None: + with pytest.warns(DeprecationWarning, match="Using the current directory as output_dir is deprecated"): + async with JupyterCodeExecutor(output_dir=".") as executor: + _ = executor.output_dir + + +@pytest.mark.asyncio +async def test_runtime_error_not_started() -> None: + executor = JupyterCodeExecutor() + code_blocks = [CodeBlock(code="print('hello world!')", language="python")] + with pytest.raises(RuntimeError, match="Executor must be started before executing cells"): + await executor.execute_code_blocks(code_blocks, CancellationToken()) diff --git a/python/packages/autogen-ext/tests/code_executors/test_user_defined_functions.py b/python/packages/autogen-ext/tests/code_executors/test_user_defined_functions.py index c8dd6eb5fad1..ebfdf53287ab 100644 --- a/python/packages/autogen-ext/tests/code_executors/test_user_defined_functions.py +++ b/python/packages/autogen-ext/tests/code_executors/test_user_defined_functions.py @@ -3,6 +3,7 @@ import os import tempfile +from pathlib import Path import polars import pytest @@ -59,6 +60,7 @@ async def test_can_load_function_with_reqs() -> None: with tempfile.TemporaryDirectory() as temp_dir: cancellation_token = CancellationToken() executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[load_data]) + await executor.start() code = f"""from {executor.functions_module} import load_data import polars @@ -75,19 +77,24 @@ async def test_can_load_function_with_reqs() -> None: assert result.output == f"John{os.linesep}" assert result.exit_code == 0 + await executor.stop() -def test_local_formatted_prompt() -> None: + +async def test_local_formatted_prompt() -> None: assert_str = '''def add_two_numbers(a: int, b: int) -> int: """Add two numbers together.""" ''' with tempfile.TemporaryDirectory() as temp_dir: executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[add_two_numbers]) + await executor.start() result = executor.format_functions_for_prompt() assert assert_str in result + await executor.stop() + -def test_local_formatted_prompt_str_func() -> None: +async def test_local_formatted_prompt_str_func() -> None: func = FunctionWithRequirements.from_str( ''' def add_two_numbers(a: int, b: int) -> int: @@ -102,16 +109,20 @@ def add_two_numbers(a: int, b: int) -> int: with tempfile.TemporaryDirectory() as temp_dir: executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[func]) + await executor.start() result = executor.format_functions_for_prompt() assert assert_str in result + await executor.stop() + @pytest.mark.asyncio async def test_can_load_function() -> None: with tempfile.TemporaryDirectory() as temp_dir: cancellation_token = CancellationToken() executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[add_two_numbers]) + await executor.start() code = f"""from {executor.functions_module} import add_two_numbers print(add_two_numbers(1, 2))""" @@ -124,12 +135,17 @@ async def test_can_load_function() -> None: assert result.output == f"3{os.linesep}" assert result.exit_code == 0 + await executor.stop() + @pytest.mark.asyncio async def test_fails_for_function_incorrect_import() -> None: with tempfile.TemporaryDirectory() as temp_dir: cancellation_token = CancellationToken() executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[function_incorrect_import]) + + await executor.start() + code = f"""from {executor.functions_module} import function_incorrect_import function_incorrect_import()""" @@ -141,12 +157,17 @@ async def test_fails_for_function_incorrect_import() -> None: cancellation_token=cancellation_token, ) + await executor.stop() + @pytest.mark.asyncio async def test_fails_for_function_incorrect_dep() -> None: with tempfile.TemporaryDirectory() as temp_dir: cancellation_token = CancellationToken() executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[function_incorrect_dep]) + + await executor.start() + code = f"""from {executor.functions_module} import function_incorrect_dep function_incorrect_dep()""" @@ -158,6 +179,8 @@ async def test_fails_for_function_incorrect_dep() -> None: cancellation_token=cancellation_token, ) + await executor.stop() + @pytest.mark.asyncio async def test_can_load_str_function_with_reqs() -> None: @@ -172,6 +195,8 @@ def add_two_numbers(a: int, b: int) -> int: cancellation_token = CancellationToken() executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[func]) + await executor.start() + code = f"""from {executor.functions_module} import add_two_numbers print(add_two_numbers(1, 2))""" @@ -184,6 +209,8 @@ def add_two_numbers(a: int, b: int) -> int: assert result.output == f"3{os.linesep}" assert result.exit_code == 0 + await executor.stop() + def test_cant_load_broken_str_function_with_reqs() -> None: with pytest.raises(ValueError): @@ -209,6 +236,8 @@ def add_two_numbers(a: int, b: int) -> int: cancellation_token = CancellationToken() executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[func]) + await executor.start() + code = f"""from {executor.functions_module} import add_two_numbers print(add_two_numbers(object(), False))""" @@ -220,3 +249,71 @@ def add_two_numbers(a: int, b: int) -> int: ) assert "TypeError: unsupported operand type(s) for +:" in result.output assert result.exit_code == 1 + + await executor.stop() + + +@pytest.mark.asyncio +async def test_error_wrong_path() -> None: + with tempfile.TemporaryDirectory() as temp_dir: + executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[]) + await executor.start() + + code_blocks = [ + CodeBlock( + code="""with open("/nonexistent_dir/test.txt", "w") as f: + f.write("hello word")""", + language="python", + ) + ] + + result = await executor.execute_code_blocks(code_blocks, CancellationToken()) + assert result.exit_code != 0 + assert "No such file or directory" in result.output + + await executor.stop() + + +@pytest.mark.asyncio +async def test_deprecated_warning() -> None: + with pytest.warns(DeprecationWarning, match="Using the current directory as work_dir is deprecated."): + executor = LocalCommandLineCodeExecutor(work_dir=".", functions=[]) + await executor.start() + + code_block = CodeBlock(code='echo "hello word"', language="sh") + result = await executor.execute_code_blocks([code_block], CancellationToken()) + + assert result.exit_code == 0 + assert "hello word" in result.output + + await executor.stop() + + +@pytest.mark.asyncio +async def test_default_work_dir_is_temp() -> None: + executor = LocalCommandLineCodeExecutor(functions=[]) + await executor.start() + + assert executor.work_dir != Path(".") + + system_temp = tempfile.gettempdir() + assert system_temp in str(executor.work_dir) + + await executor.stop() + + +def test_invalid_timeout() -> None: + with pytest.raises(ValueError, match="Timeout must be greater than or equal to 1."): + _ = LocalCommandLineCodeExecutor(timeout=0) + + +def test_python_identifier() -> None: + with pytest.raises(ValueError, match="Module name must be a valid Python identifier"): + # Using a name with an hyphen is an example of an invalid Python identifier + _ = LocalCommandLineCodeExecutor(functions_module="invalid-identifier") + + +@pytest.mark.asyncio +async def test_create_temp_dir() -> None: + executor = LocalCommandLineCodeExecutor() + assert executor.work_dir.is_dir() diff --git a/python/packages/autogen-ext/tests/mcp_server_comprehensive.py b/python/packages/autogen-ext/tests/mcp_server_comprehensive.py new file mode 100644 index 000000000000..7a7615e080ea --- /dev/null +++ b/python/packages/autogen-ext/tests/mcp_server_comprehensive.py @@ -0,0 +1,248 @@ +import asyncio +import json +import logging +from datetime import datetime +from typing import Any, Dict, Optional + +from mcp import PromptsCapability, ResourcesCapability, ServerCapabilities, ToolsCapability +from mcp.server import Server +from mcp.server.models import InitializationOptions +from mcp.server.stdio import stdio_server +from mcp.types import ( + GetPromptResult, + Prompt, + PromptArgument, + PromptMessage, + Resource, + TextContent, + Tool, +) +from pydantic import AnyUrl + +# Configure logging +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +# Sample data for demonstration +SAMPLE_DATA = { + "users": [ + {"id": 1, "name": "Alice", "email": "alice@example.com", "department": "Engineering"}, + {"id": 2, "name": "Bob", "email": "bob@example.com", "department": "Sales"}, + {"id": 3, "name": "Charlie", "email": "charlie@example.com", "department": "Marketing"}, + ], + "projects": [ + {"id": 1, "name": "Project Alpha", "status": "active", "team_size": 5}, + {"id": 2, "name": "Project Beta", "status": "completed", "team_size": 3}, + {"id": 3, "name": "Project Gamma", "status": "planning", "team_size": 2}, + ], +} + + +class SimpleMcpServer: + """A simple MCP server demonstrating basic functionality.""" + + def __init__(self) -> None: + self.server: Server[object] = Server("simple-mcp-server") + self.register_handlers() # type: ignore[no-untyped-call] + + def register_handlers(self) -> None: + """Register all MCP handlers.""" + + # Prompts + @self.server.list_prompts() # type: ignore[no-untyped-call,misc] + async def list_prompts() -> list[Prompt]: # pyright: ignore[reportUnusedFunction] + """List available prompts.""" + return [ + Prompt( + name="code_review", + description="Generate a comprehensive code review for a given piece of code", + arguments=[ + PromptArgument( + name="code", + description="The code to review", + required=True, + ), + PromptArgument( + name="language", + description="Programming language of the code", + required=True, + ), + ], + ), + Prompt( + name="documentation", + description="Generate documentation for code or APIs", + arguments=[ + PromptArgument( + name="content", + description="The content to document", + required=True, + ), + ], + ), + ] + + @self.server.get_prompt() # type: ignore[no-untyped-call,misc] + async def get_prompt(name: str, arguments: Optional[Dict[str, str]] = None) -> GetPromptResult: # pyright: ignore[reportUnusedFunction] + """Get a specific prompt with arguments.""" + if not arguments: + arguments = {} + + if name == "code_review": + code = arguments.get("code", "// No code provided") + language = arguments.get("language", "unknown") + + return GetPromptResult( + description=f"Code review for {language} code", + messages=[ + PromptMessage( + role="user", + content=TextContent( + type="text", + text=f"Please review this {language} code:\n\n```{language}\n{code}\n```", + ), + ), + ], + ) + + elif name == "documentation": + content = arguments.get("content", "No content provided") + + return GetPromptResult( + description="Documentation generation", + messages=[ + PromptMessage( + role="user", + content=TextContent( + type="text", + text=f"Please generate documentation for:\n\n{content}", + ), + ), + ], + ) + + else: + raise ValueError(f"Unknown prompt: {name}") + + # Resources + @self.server.list_resources() # type: ignore[no-untyped-call,misc] + async def list_resources() -> list[Resource]: # pyright: ignore[reportUnusedFunction] + """List available resources.""" + return [ + Resource( + uri=AnyUrl("file:///company/users.json"), + name="Company Users", + description="List of all company users", + mimeType="application/json", + ), + Resource( + uri=AnyUrl("file:///company/projects.json"), + name="Active Projects", + description="Current projects", + mimeType="application/json", + ), + ] + + @self.server.read_resource() # type: ignore[no-untyped-call,misc] + async def read_resource(uri: AnyUrl) -> str: # pyright: ignore[reportUnusedFunction] + """Read a specific resource.""" + uri_str = str(uri) + + if uri_str == "file:///company/users.json": + return json.dumps(SAMPLE_DATA["users"], indent=2) + + elif uri_str == "file:///company/projects.json": + return json.dumps(SAMPLE_DATA["projects"], indent=2) + + else: + raise ValueError(f"Unknown resource: {uri_str}") + + # Tools + @self.server.list_tools() # type: ignore[no-untyped-call,misc] + async def list_tools() -> list[Tool]: # pyright: ignore[reportUnusedFunction] + """List available tools.""" + return [ + Tool( + name="echo", + description="Echo back the input text", + inputSchema={ + "type": "object", + "properties": { + "text": { + "type": "string", + "description": "Text to echo back", + } + }, + "required": ["text"], + }, + ), + Tool( + name="get_time", + description="Get the current time", + inputSchema={ + "type": "object", + "properties": {}, + }, + ), + ] + + @self.server.call_tool() # type: ignore[no-untyped-call,misc] + async def call_tool(name: str, arguments: Optional[Dict[str, Any]] = None) -> list[TextContent]: # pyright: ignore[reportUnusedFunction] + """Call a specific tool.""" + if not arguments: + arguments = {} + + if name == "echo": + text = arguments.get("text", "") + return [ + TextContent( + type="text", + text=f"Echo: {text}", + ) + ] + + elif name == "get_time": + current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + return [ + TextContent( + type="text", + text=f"Current time: {current_time}", + ) + ] + + else: + raise ValueError(f"Unknown tool: {name}") + + async def run(self) -> None: + """Run the MCP server.""" + # Server capabilities + init_options = InitializationOptions( + server_name="simple-mcp-server", + server_version="1.0.0", + capabilities=ServerCapabilities( + prompts=PromptsCapability(listChanged=True), + resources=ResourcesCapability( + subscribe=True, + listChanged=True, + ), + tools=ToolsCapability(listChanged=True), + ), + ) + + # Run the server + async with stdio_server() as (read_stream, write_stream): + await self.server.run( + read_stream, + write_stream, + init_options, + ) + + +async def main() -> None: + """Main entry point.""" + server: SimpleMcpServer = SimpleMcpServer() + await server.run() # type: ignore[no-untyped-call] + + +if __name__ == "__main__": + asyncio.run(main()) # type: ignore[no-untyped-call] diff --git a/python/packages/autogen-ext/tests/memory/test_chroma_memory.py b/python/packages/autogen-ext/tests/memory/test_chroma_memory.py index 163b77f2c257..f62c91c7dde1 100644 --- a/python/packages/autogen-ext/tests/memory/test_chroma_memory.py +++ b/python/packages/autogen-ext/tests/memory/test_chroma_memory.py @@ -4,7 +4,21 @@ from autogen_core.memory import MemoryContent, MemoryMimeType from autogen_core.model_context import BufferedChatCompletionContext from autogen_core.models import UserMessage -from autogen_ext.memory.chromadb import ChromaDBVectorMemory, PersistentChromaDBVectorMemoryConfig +from autogen_ext.memory.chromadb import ( + ChromaDBVectorMemory, + CustomEmbeddingFunctionConfig, + DefaultEmbeddingFunctionConfig, + HttpChromaDBVectorMemoryConfig, + OpenAIEmbeddingFunctionConfig, + PersistentChromaDBVectorMemoryConfig, + SentenceTransformerEmbeddingFunctionConfig, +) + +# Skip all tests if ChromaDB is not available +try: + import chromadb # pyright: ignore[reportUnusedImport] +except ImportError: + pytest.skip("ChromaDB not available", allow_module_level=True) @pytest.fixture @@ -240,3 +254,189 @@ async def test_component_serialization(base_config: PersistentChromaDBVectorMemo await memory.close() await loaded_memory.close() + + +@pytest.mark.asyncio +def test_http_config(tmp_path: Path) -> None: + """Test HTTP ChromaDB configuration.""" + config = HttpChromaDBVectorMemoryConfig( + collection_name="test_http", + host="localhost", + port=8000, + ssl=False, + headers={"Authorization": "Bearer test-token"}, + ) + + assert config.client_type == "http" + assert config.host == "localhost" + assert config.port == 8000 + assert config.ssl is False + assert config.headers == {"Authorization": "Bearer test-token"} + + +# ============================================================================ +# Embedding Function Configuration Tests +# ============================================================================ + + +@pytest.mark.asyncio +async def test_default_embedding_function(tmp_path: Path) -> None: + """Test ChromaDB memory with default embedding function.""" + config = PersistentChromaDBVectorMemoryConfig( + collection_name="test_default_embedding", + allow_reset=True, + persistence_path=str(tmp_path / "chroma_db_default"), + embedding_function_config=DefaultEmbeddingFunctionConfig(), + ) + + memory = ChromaDBVectorMemory(config=config) + await memory.clear() + + # Add test content + await memory.add( + MemoryContent( + content="Default embedding function test content", + mime_type=MemoryMimeType.TEXT, + metadata={"test": "default_embedding"}, + ) + ) + + # Query and verify + results = await memory.query("default embedding test") + assert len(results.results) > 0 + assert any("Default embedding" in str(r.content) for r in results.results) + + await memory.close() + + +@pytest.mark.asyncio +async def test_sentence_transformer_embedding_function(tmp_path: Path) -> None: + """Test ChromaDB memory with SentenceTransformer embedding function.""" + config = PersistentChromaDBVectorMemoryConfig( + collection_name="test_st_embedding", + allow_reset=True, + persistence_path=str(tmp_path / "chroma_db_st"), + embedding_function_config=SentenceTransformerEmbeddingFunctionConfig( + model_name="all-MiniLM-L6-v2" # Use default model for testing + ), + ) + + memory = ChromaDBVectorMemory(config=config) + await memory.clear() + + # Add test content + await memory.add( + MemoryContent( + content="SentenceTransformer embedding function test content", + mime_type=MemoryMimeType.TEXT, + metadata={"test": "sentence_transformer"}, + ) + ) + + # Query and verify + results = await memory.query("SentenceTransformer embedding test") + assert len(results.results) > 0 + assert any("SentenceTransformer" in str(r.content) for r in results.results) + + await memory.close() + + +@pytest.mark.asyncio +async def test_custom_embedding_function(tmp_path: Path) -> None: + """Test ChromaDB memory with custom embedding function.""" + from collections.abc import Sequence + + class MockEmbeddingFunction: + def __call__(self, input: Sequence[str]) -> list[list[float]]: + # Return a batch of embeddings (list of lists) + return [[0.0] * 384 for _ in input] + + config = PersistentChromaDBVectorMemoryConfig( + collection_name="test_custom_embedding", + allow_reset=True, + persistence_path=str(tmp_path / "chroma_db_custom"), + embedding_function_config=CustomEmbeddingFunctionConfig(function=MockEmbeddingFunction, params={}), + ) + memory = ChromaDBVectorMemory(config=config) + await memory.clear() + await memory.add( + MemoryContent( + content="Custom embedding function test content", + mime_type=MemoryMimeType.TEXT, + metadata={"test": "custom_embedding"}, + ) + ) + results = await memory.query("custom embedding test") + assert len(results.results) > 0 + assert any("Custom embedding" in str(r.content) for r in results.results) + await memory.close() + + +@pytest.mark.asyncio +async def test_openai_embedding_function(tmp_path: Path) -> None: + """Test OpenAI embedding function configuration (without actual API call).""" + config = PersistentChromaDBVectorMemoryConfig( + collection_name="test_openai_embedding", + allow_reset=True, + persistence_path=str(tmp_path / "chroma_db_openai"), + embedding_function_config=OpenAIEmbeddingFunctionConfig( + api_key="test-key", model_name="text-embedding-3-small" + ), + ) + + # Just test that the config is valid - don't actually try to use OpenAI API + assert config.embedding_function_config.function_type == "openai" + assert config.embedding_function_config.api_key == "test-key" + assert config.embedding_function_config.model_name == "text-embedding-3-small" + + +@pytest.mark.asyncio +async def test_embedding_function_error_handling(tmp_path: Path) -> None: + """Test error handling for embedding function configurations.""" + + def failing_embedding_function() -> None: + """A function that raises an error.""" + raise ValueError("Test embedding function error") + + config = PersistentChromaDBVectorMemoryConfig( + collection_name="test_error_embedding", + allow_reset=True, + persistence_path=str(tmp_path / "chroma_db_error"), + embedding_function_config=CustomEmbeddingFunctionConfig(function=failing_embedding_function, params={}), + ) + + memory = ChromaDBVectorMemory(config=config) + + # Should raise an error when trying to initialize + with pytest.raises((ValueError, Exception)): # Catch ValueError or any other exception + await memory.add(MemoryContent(content="This should fail", mime_type=MemoryMimeType.TEXT)) + + await memory.close() + + +def test_embedding_function_config_validation() -> None: + """Test validation of embedding function configurations.""" + + # Test default config + default_config = DefaultEmbeddingFunctionConfig() + assert default_config.function_type == "default" + + # Test SentenceTransformer config + st_config = SentenceTransformerEmbeddingFunctionConfig(model_name="test-model") + assert st_config.function_type == "sentence_transformer" + assert st_config.model_name == "test-model" + + # Test OpenAI config + openai_config = OpenAIEmbeddingFunctionConfig(api_key="test-key", model_name="test-model") + assert openai_config.function_type == "openai" + assert openai_config.api_key == "test-key" + assert openai_config.model_name == "test-model" + + # Test custom config + def dummy_function() -> None: + return None + + custom_config = CustomEmbeddingFunctionConfig(function=dummy_function, params={"test": "value"}) + assert custom_config.function_type == "custom" + assert custom_config.function == dummy_function + assert custom_config.params == {"test": "value"} diff --git a/python/packages/autogen-ext/tests/memory/test_mem0.py b/python/packages/autogen-ext/tests/memory/test_mem0.py new file mode 100644 index 000000000000..27235e1305b9 --- /dev/null +++ b/python/packages/autogen-ext/tests/memory/test_mem0.py @@ -0,0 +1,530 @@ +import os +import uuid +from datetime import datetime +from typing import Any, Dict +from unittest.mock import MagicMock, patch + +import pytest +from autogen_core.memory import MemoryContent, MemoryMimeType +from autogen_core.model_context import BufferedChatCompletionContext +from autogen_core.models import SystemMessage, UserMessage +from autogen_ext.memory.mem0 import Mem0Memory, Mem0MemoryConfig +from dotenv import load_dotenv + +# Load environment variables from .env file +load_dotenv() + +# Skip tests if required environment variables are not set +mem0_api_key = os.environ.get("MEM0_API_KEY") +requires_mem0_api = pytest.mark.skipif(mem0_api_key is None, reason="MEM0_API_KEY environment variable not set") + +# Skip tests if mem0ai is not installed +mem0 = pytest.importorskip("mem0") + +# Define local configuration at the top of the module +FULL_LOCAL_CONFIG: Dict[str, Any] = { + "history_db_path": ":memory:", # Use in-memory DB for tests + "graph_store": { + "provider": "mock_graph", + "config": {"url": "mock://localhost:7687", "username": "mock", "password": "mock_password"}, + }, + "embedder": { + "provider": "mock_embedder", + "config": { + "model": "mock-embedding-model", + "embedding_dims": 1024, + "api_key": "mock-api-key", + }, + }, + "vector_store": {"provider": "mock_vector", "config": {"path": ":memory:", "collection_name": "test_memories"}}, + "llm": { + "provider": "mock_llm", + "config": { + "model": "mock-chat-model", + "api_key": "mock-api-key", + }, + }, +} + + +@pytest.fixture +def full_local_config() -> Dict[str, Any]: + """Return the local configuration for testing.""" + return FULL_LOCAL_CONFIG + + +@pytest.fixture +def cloud_config() -> Mem0MemoryConfig: + """Create cloud configuration with real API key.""" + api_key = os.environ.get("MEM0_API_KEY") + return Mem0MemoryConfig(user_id="test-user", limit=3, is_cloud=True, api_key=api_key) + + +@pytest.fixture +def local_config() -> Mem0MemoryConfig: + """Create local configuration for testing.""" + return Mem0MemoryConfig(user_id="test-user", limit=3, is_cloud=False, config={"path": ":memory:"}) + + +@pytest.mark.asyncio +@patch("autogen_ext.memory.mem0._mem0.Memory0") +async def test_basic_workflow(mock_mem0_class: MagicMock, local_config: Mem0MemoryConfig) -> None: + """Test basic memory operations.""" + # Setup mock + mock_mem0 = MagicMock() + mock_mem0_class.from_config.return_value = mock_mem0 + + # Mock search results + mock_mem0.search.return_value = [ + { + "memory": "Paris is known for the Eiffel Tower and amazing cuisine.", + "score": 0.95, + "metadata": {"category": "city", "country": "France"}, + } + ] + + memory = Mem0Memory( + user_id=local_config.user_id, + limit=local_config.limit, + is_cloud=local_config.is_cloud, + api_key=local_config.api_key, + config=local_config.config, + ) + + # Add content to memory + await memory.add( + MemoryContent( + content="Paris is known for the Eiffel Tower and amazing cuisine.", + mime_type=MemoryMimeType.TEXT, + metadata={"category": "city", "country": "France"}, + ) + ) + + # Verify add was called correctly + mock_mem0.add.assert_called_once() + call_args = mock_mem0.add.call_args[0] + + # Extract content from the list of dict structure: [{'content': '...', 'role': 'user'}] + actual_content = call_args[0][0]["content"] + assert actual_content == "Paris is known for the Eiffel Tower and amazing cuisine." + + call_kwargs = mock_mem0.add.call_args[1] + assert call_kwargs["metadata"] == {"category": "city", "country": "France"} + + # Query memory + results = await memory.query("Tell me about Paris") + + # Verify search was called correctly + mock_mem0.search.assert_called_once() + search_args = mock_mem0.search.call_args + assert search_args[0][0] == "Tell me about Paris" + assert search_args[1]["user_id"] == "test-user" + assert search_args[1]["limit"] == 3 + + # Verify results + assert len(results.results) == 1 + assert "Paris" in str(results.results[0].content) + res_metadata = results.results[0].metadata + assert res_metadata is not None and res_metadata.get("score") == 0.95 + assert res_metadata is not None and res_metadata.get("country") == "France" + + # Test clear (only do this once) + await memory.clear() + mock_mem0.delete_all.assert_called_once_with(user_id="test-user") + + # Cleanup + await memory.close() + + +@requires_mem0_api +@pytest.mark.asyncio +@patch("autogen_ext.memory.mem0.MemoryClient") # Patch MemoryClient instead of Memory0 +async def test_basic_workflow_with_cloud(mock_memory_client_class: MagicMock, cloud_config: Mem0MemoryConfig) -> None: + """Test basic memory operations with cloud client (mocked instead of real API).""" + # Setup mock + mock_client = MagicMock() + mock_memory_client_class.return_value = mock_client + + # Mock search results + mock_client.search.return_value = [ + { + "memory": "Test memory content for cloud", + "score": 0.98, + "metadata": {"test": True, "source": "cloud"}, + } + ] + + memory = Mem0Memory( + user_id=cloud_config.user_id, + limit=cloud_config.limit, + is_cloud=cloud_config.is_cloud, + api_key=cloud_config.api_key, + config=cloud_config.config, + ) + + # Generate a unique test content string + test_content = f"Test memory content {uuid.uuid4()}" + + # Add content to memory + await memory.add( + MemoryContent( + content=test_content, + mime_type=MemoryMimeType.TEXT, + metadata={"test": True, "timestamp": datetime.now().isoformat()}, + ) + ) + + # Verify add was called correctly + mock_client.add.assert_called_once() + call_args = mock_client.add.call_args + + # Extract content from list of dict structure: [{'content': '...', 'role': 'user'}] + actual_content = call_args[0][0][0]["content"] # call_args[0][0] gets the first positional arg (the list) + assert test_content in actual_content + + assert call_args[1]["user_id"] == cloud_config.user_id + assert call_args[1]["metadata"]["test"] is True + + # Query memory + results = await memory.query(test_content) + + # Verify search was called correctly + mock_client.search.assert_called_once() + search_args = mock_client.search.call_args + assert test_content in search_args[0][0] + assert search_args[1]["user_id"] == cloud_config.user_id + + # Verify results + assert len(results.results) == 1 + assert "Test memory content for cloud" in str(results.results[0].content) + assert results.results[0].metadata is not None + assert results.results[0].metadata.get("score") == 0.98 + + # Test clear + await memory.clear() + mock_client.delete_all.assert_called_once_with(user_id=cloud_config.user_id) + + # Cleanup + await memory.close() + + +@pytest.mark.asyncio +@patch("autogen_ext.memory.mem0._mem0.Memory0") +async def test_metadata_handling(mock_mem0_class: MagicMock, local_config: Mem0MemoryConfig) -> None: + """Test metadata handling.""" + # Setup mock + mock_mem0 = MagicMock() + mock_mem0_class.from_config.return_value = mock_mem0 + + # Setup mock search results with rich metadata + mock_mem0.search.return_value = [ + { + "memory": "Test content with metadata", + "score": 0.9, + "metadata": {"test_category": "test", "test_priority": 1, "test_weight": 0.5, "test_verified": True}, + "created_at": "2023-01-01T12:00:00", + "updated_at": "2023-01-02T12:00:00", + "categories": ["test", "example"], + } + ] + + memory = Mem0Memory( + user_id=local_config.user_id, + limit=local_config.limit, + is_cloud=local_config.is_cloud, + api_key=local_config.api_key, + config=local_config.config, + ) + + # Add content with metadata + test_content = "Test content with specific metadata" + content = MemoryContent( + content=test_content, + mime_type=MemoryMimeType.TEXT, + metadata={"test_category": "test", "test_priority": 1, "test_weight": 0.5, "test_verified": True}, + ) + await memory.add(content) + + # Verify metadata was passed correctly + add_kwargs = mock_mem0.add.call_args[1] + assert add_kwargs["metadata"]["test_category"] == "test" + assert add_kwargs["metadata"]["test_priority"] == 1 + assert add_kwargs["metadata"]["test_weight"] == 0.5 + assert add_kwargs["metadata"]["test_verified"] is True + + # Query and check returned metadata + results = await memory.query(test_content) + assert len(results.results) == 1 + result = results.results[0] + + # Verify metadata in results + assert result.metadata is not None and result.metadata.get("test_category") == "test" + assert result.metadata is not None and result.metadata.get("test_priority") == 1 + assert result.metadata is not None and isinstance(result.metadata.get("test_weight"), float) + assert result.metadata is not None and result.metadata.get("test_verified") is True + assert result.metadata is not None and "created_at" in result.metadata + assert result.metadata is not None and "updated_at" in result.metadata + assert result.metadata is not None and result.metadata.get("categories") == ["test", "example"] + + # Cleanup + await memory.close() + + +@pytest.mark.asyncio +@patch("autogen_ext.memory.mem0._mem0.Memory0") +async def test_update_context(mock_mem0_class: MagicMock, local_config: Mem0MemoryConfig) -> None: + """Test updating model context with retrieved memories.""" + # Setup mock + mock_mem0 = MagicMock() + mock_mem0_class.from_config.return_value = mock_mem0 + + # Setup mock search results + mock_mem0.search.return_value = [ + {"memory": "Jupiter is the largest planet in our solar system.", "score": 0.9}, + {"memory": "Jupiter has many moons, including Ganymede, Europa, and Io.", "score": 0.8}, + ] + + memory = Mem0Memory( + user_id=local_config.user_id, + limit=local_config.limit, + is_cloud=local_config.is_cloud, + api_key=local_config.api_key, + config=local_config.config, + ) + + # Create a model context with a message + context = BufferedChatCompletionContext(buffer_size=5) + await context.add_message(UserMessage(content="Tell me about Jupiter", source="user")) + + # Update context with memory + result = await memory.update_context(context) + + # Verify results + assert len(result.memories.results) == 2 + assert "Jupiter" in str(result.memories.results[0].content) + + # Verify search was called with correct query + mock_mem0.search.assert_called_once() + search_args = mock_mem0.search.call_args + assert "Jupiter" in search_args[0][0] + + # Verify context was updated with a system message + messages = await context.get_messages() + assert len(messages) == 2 # Original message + system message with memories + + # Verify system message content + system_message = messages[1] + assert isinstance(system_message, SystemMessage) + assert "Jupiter is the largest planet" in system_message.content + assert "Jupiter has many moons" in system_message.content + + # Cleanup + await memory.close() + + +@pytest.mark.asyncio +@patch("autogen_ext.memory.mem0._mem0.MemoryClient") # Patch for cloud mode +async def test_component_serialization(mock_memory_client_class: MagicMock) -> None: + """Test serialization and deserialization of the component.""" + # Setup mock + mock_client = MagicMock() + mock_memory_client_class.return_value = mock_client + + # Create configuration + user_id = str(uuid.uuid4()) + config = Mem0MemoryConfig( + user_id=user_id, + limit=5, + is_cloud=True, + ) + + # Create memory instance + memory = Mem0Memory( + user_id=config.user_id, + limit=config.limit, + is_cloud=config.is_cloud, + api_key=config.api_key, + config=config.config, + ) + + # Dump config + memory_config = memory.dump_component() + + # Verify dumped config + assert memory_config.config["user_id"] == user_id + assert memory_config.config["limit"] == 5 + assert memory_config.config["is_cloud"] is True + + # Load from config + loaded_memory = Mem0Memory( + user_id=config.user_id, + limit=config.limit, + is_cloud=config.is_cloud, + api_key=config.api_key, + config=config.config, + ) + + # Verify loaded instance + assert isinstance(loaded_memory, Mem0Memory) + assert loaded_memory.user_id == user_id + assert loaded_memory.limit == 5 + assert loaded_memory.is_cloud is True + assert loaded_memory.config is None + + # Cleanup + await memory.close() + await loaded_memory.close() + + +@pytest.mark.asyncio +@patch("autogen_ext.memory.mem0._mem0.Memory0") +async def test_result_format_handling(mock_mem0_class: MagicMock, local_config: Mem0MemoryConfig) -> None: + """Test handling of different result formats.""" + # Setup mock + mock_mem0 = MagicMock() + mock_mem0_class.from_config.return_value = mock_mem0 + + # Test dictionary format with "results" key + mock_mem0.search.return_value = { + "results": [ + {"memory": "Dictionary format result 1", "score": 0.9}, + {"memory": "Dictionary format result 2", "score": 0.8}, + ] + } + + memory = Mem0Memory( + user_id=local_config.user_id, + limit=local_config.limit, + is_cloud=local_config.is_cloud, + api_key=local_config.api_key, + config=local_config.config, + ) + + # Query with dictionary format + results_dict = await memory.query("test query") + + # Verify results were extracted correctly + assert len(results_dict.results) == 2 + assert "Dictionary format result 1" in str(results_dict.results[0].content) + + # Test list format + mock_mem0.search.return_value = [ + {"memory": "List format result 1", "score": 0.9}, + {"memory": "List format result 2", "score": 0.8}, + ] + + # Query with list format + results_list = await memory.query("test query") + + # Verify results were processed correctly + assert len(results_list.results) == 2 + assert "List format result 1" in str(results_list.results[0].content) + + # Cleanup + await memory.close() + + +@pytest.mark.asyncio +@patch("autogen_ext.memory.mem0._mem0.Memory0") +async def test_init_with_local_config(mock_mem0_class: MagicMock, full_local_config: Dict[str, Any]) -> None: + """Test initializing memory with local configuration.""" + # Setup mock + mock_mem0 = MagicMock() + mock_mem0_class.from_config.return_value = mock_mem0 + + # Initialize memory with local config + memory = Mem0Memory(user_id="test-local-config-user", limit=10, is_cloud=False, config=full_local_config) + + # Verify configuration was passed correctly + mock_mem0_class.from_config.assert_called_once() + + # Verify memory instance properties (use type: ignore or add public properties) + assert memory._user_id == "test-local-config-user" # type: ignore + assert memory._limit == 10 # type: ignore + assert memory._is_cloud is False # type: ignore + assert memory._config == full_local_config # type: ignore + + # Test serialization with local config + memory_config = memory.dump_component() + + # Verify serialized config + assert memory_config.config["user_id"] == "test-local-config-user" + assert memory_config.config["is_cloud"] is False + + # Cleanup + await memory.close() + + +@pytest.mark.asyncio +@patch("autogen_ext.memory.mem0._mem0.Memory0") # Patches the underlying mem0.Memory class +async def test_local_config_with_memory_operations( + mock_mem0_class: MagicMock, + full_local_config: Dict[str, Any], # full_local_config fixture provides the mock config +) -> None: + """Test memory operations with local configuration.""" + # Setup mock for the instance that will be created by Mem0Memory + mock_mem0_instance = MagicMock() + mock_mem0_class.from_config.return_value = mock_mem0_instance + + # Mock search results from the mem0 instance + mock_mem0_instance.search.return_value = [ + { + "memory": "Test local config memory content", + "score": 0.92, + "metadata": {"config_type": "local", "test_case": "advanced"}, + } + ] + + # Initialize Mem0Memory with is_cloud=False and the full_local_config + memory = Mem0Memory(user_id="test-local-config-user", limit=10, is_cloud=False, config=full_local_config) + + # Verify that mem0.Memory.from_config was called with the provided config + mock_mem0_class.from_config.assert_called_once_with(config_dict=full_local_config) + + # Add memory content + test_content_str = "Testing local configuration memory operations" + await memory.add( + MemoryContent( + content=test_content_str, + mime_type=MemoryMimeType.TEXT, + metadata={"config_type": "local", "test_case": "advanced"}, + ) + ) + + # Verify add was called on the mock_mem0_instance + mock_mem0_instance.add.assert_called_once() + + # Query memory + results = await memory.query("local configuration test") + + # Verify search was called on the mock_mem0_instance + mock_mem0_instance.search.assert_called_once_with( + "local configuration test", user_id="test-local-config-user", limit=10 + ) + + # Verify results + assert len(results.results) == 1 + assert "Test local config memory content" in str(results.results[0].content) + res_metadata = results.results[0].metadata + assert res_metadata is not None and res_metadata.get("score") == 0.92 + assert res_metadata is not None and res_metadata.get("config_type") == "local" + + # Test serialization with local config + memory_config = memory.dump_component() + + # Verify serialized config + assert memory_config.config["user_id"] == "test-local-config-user" + assert memory_config.config["is_cloud"] is False + assert "config" in memory_config.config + assert memory_config.config["config"]["history_db_path"] == ":memory:" + + # Test clear + await memory.clear() + mock_mem0_instance.delete_all.assert_called_once_with(user_id="test-local-config-user") + + # Cleanup + await memory.close() + + +if __name__ == "__main__": + pytest.main(["-xvs", __file__]) diff --git a/python/packages/autogen-ext/tests/memory/test_redis_memory.py b/python/packages/autogen-ext/tests/memory/test_redis_memory.py new file mode 100644 index 000000000000..2ab25f55c8f3 --- /dev/null +++ b/python/packages/autogen-ext/tests/memory/test_redis_memory.py @@ -0,0 +1,421 @@ +from collections.abc import AsyncGenerator +from unittest.mock import MagicMock, patch + +import pytest +import pytest_asyncio +from autogen_core.memory import MemoryContent, MemoryMimeType +from autogen_core.model_context import BufferedChatCompletionContext +from autogen_core.models import UserMessage +from autogen_ext.memory.redis import RedisMemory, RedisMemoryConfig +from pydantic import ValidationError +from redis import Redis +from redisvl.exceptions import RedisSearchError + + +@pytest.mark.asyncio +async def test_redis_memory_add_with_mock() -> None: + with patch("autogen_ext.memory.redis._redis_memory.SemanticMessageHistory") as MockHistory: + mock_history = MagicMock() + MockHistory.return_value = mock_history + + config = RedisMemoryConfig() + memory = RedisMemory(config=config) + + content = MemoryContent(content="test content", mime_type=MemoryMimeType.TEXT, metadata={"foo": "bar"}) + await memory.add(content) + mock_history.add_message.assert_called_once() + + +@pytest.mark.asyncio +async def test_redis_memory_query_with_mock() -> None: + with patch("autogen_ext.memory.redis._redis_memory.SemanticMessageHistory") as MockHistory: + mock_history = MagicMock() + MockHistory.return_value = mock_history + + config = RedisMemoryConfig() + memory = RedisMemory(config=config) + + mock_history.get_relevant.return_value = [ + {"content": "test content", "tool_call_id": '{"foo": "bar", "mime_type": "text/plain"}'} + ] + result = await memory.query("test") + assert len(result.results) == 1 + assert result.results[0].content == "test content" + assert result.results[0].metadata == {"foo": "bar"} + mock_history.get_relevant.assert_called_once() + + +@pytest.mark.asyncio +async def test_redis_memory_clear_with_mock() -> None: + with patch("autogen_ext.memory.redis._redis_memory.SemanticMessageHistory") as MockHistory: + mock_history = MagicMock() + MockHistory.return_value = mock_history + + config = RedisMemoryConfig() + memory = RedisMemory(config=config) + + await memory.clear() + mock_history.clear.assert_called_once() + + +@pytest.mark.asyncio +async def test_redis_memory_close_with_mock() -> None: + with patch("autogen_ext.memory.redis._redis_memory.SemanticMessageHistory") as MockHistory: + mock_history = MagicMock() + MockHistory.return_value = mock_history + + config = RedisMemoryConfig() + memory = RedisMemory(config=config) + + await memory.close() + mock_history.delete.assert_called_once() + + +def redis_available() -> bool: + try: + client = Redis.from_url("redis://localhost:6379") # type: ignore[reportUnkownMemberType] + client.ping() # type: ignore[reportUnkownMemberType] + return True + except Exception: + return False + + +@pytest.fixture +def semantic_config() -> RedisMemoryConfig: + """Create base configuration using semantic memory.""" + return RedisMemoryConfig(top_k=5, distance_threshold=0.5, model_name="sentence-transformers/all-mpnet-base-v2") + + +@pytest_asyncio.fixture # type: ignore[reportUntypedFunctionDecorator] +async def semantic_memory(semantic_config: RedisMemoryConfig) -> AsyncGenerator[RedisMemory]: + memory = RedisMemory(semantic_config) + yield memory + await memory.close() + + +## UNIT TESTS ## +def test_memory_config() -> None: + default_config = RedisMemoryConfig() + assert default_config.redis_url == "redis://localhost:6379" + assert default_config.index_name == "chat_history" + assert default_config.prefix == "memory" + assert default_config.distance_metric == "cosine" + assert default_config.algorithm == "flat" + assert default_config.top_k == 10 + assert default_config.distance_threshold == 0.7 + assert default_config.model_name == "sentence-transformers/all-mpnet-base-v2" + + # test we can specify each of these values + url = "rediss://localhost:7010" + name = "custom name" + prefix = "custom prefix" + metric = "ip" + algorithm = "hnsw" + k = 5 + distance = 0.25 + model = "redis/langcache-embed-v1" + + custom_config = RedisMemoryConfig( + redis_url=url, + index_name=name, + prefix=prefix, + distance_metric=metric, # type: ignore[arg-type] + algorithm=algorithm, # type: ignore[arg-type] + top_k=k, + distance_threshold=distance, + model_name=model, + ) + assert custom_config.redis_url == url + assert custom_config.index_name == name + assert custom_config.prefix == prefix + assert custom_config.distance_metric == metric + assert custom_config.algorithm == algorithm + assert custom_config.top_k == k + assert custom_config.distance_threshold == distance + assert custom_config.model_name == model + + # test that Literal values are validated correctly + with pytest.raises(ValidationError): + _ = RedisMemoryConfig(distance_metric="approximate") # type: ignore[arg-type] + + with pytest.raises(ValidationError): + _ = RedisMemoryConfig(algorithm="pythagoras") # type: ignore[arg-type] + + +@pytest.mark.asyncio +@pytest.mark.skipif(not redis_available(), reason="Redis instance not available locally") +async def test_create_semantic_memory() -> None: + config = RedisMemoryConfig(index_name="semantic_agent") + memory = RedisMemory(config=config) + + assert memory.message_history is not None + await memory.close() + + +@pytest.mark.asyncio +@pytest.mark.skipif(not redis_available(), reason="Redis instance not available locally") +async def test_update_context(semantic_memory: RedisMemory) -> None: + """Test updating model context with retrieved memories.""" + await semantic_memory.clear() + + # Add content to memory + await semantic_memory.add( + MemoryContent( + content="Canada is the second largest country in the world.", + mime_type=MemoryMimeType.TEXT, + metadata={"category": "geography"}, + ) + ) + + # Create a model context with a message + context = BufferedChatCompletionContext(buffer_size=5) + await context.add_message(UserMessage(content="Tell me about Canada", source="user")) + + # Update context with memory + result = await semantic_memory.update_context(context) + + # Verify results + assert len(result.memories.results) > 0 + assert any("Canada" in str(r.content) for r in result.memories.results) + + # Verify context was updated + messages = await context.get_messages() + assert len(messages) > 1 # Should have the original message plus the memory content + + await semantic_memory.clear() + + await semantic_memory.add( + MemoryContent( + content="Napoleon was Emporor of France from 18 May 1804 to 6 April 1814.", + mime_type=MemoryMimeType.TEXT, + metadata={}, + ) + ) + await semantic_memory.add( + MemoryContent( + content="Napoleon was also Emporor during his second reign from 20 March 1815 to 22 June 1815.", + mime_type=MemoryMimeType.TEXT, + metadata={}, + ) + ) + + context = BufferedChatCompletionContext( + buffer_size=5, + initial_messages=[ + UserMessage(content="Can you tell me about the reign of Emperor Napoleon?", source="user"), + ], + ) + + updated_context = await semantic_memory.update_context(context) + assert updated_context is not None + assert updated_context.memories is not None + assert updated_context.memories.results is not None + assert len(updated_context.memories.results) == 2 + assert ( + updated_context.memories.results[0].content + == "Napoleon was Emporor of France from 18 May 1804 to 6 April 1814." + ) + assert ( + updated_context.memories.results[1].content + == "Napoleon was also Emporor during his second reign from 20 March 1815 to 22 June 1815." + ) + + +@pytest.mark.asyncio +@pytest.mark.skipif(not redis_available(), reason="Redis instance not available locally") +async def test_add_and_query(semantic_memory: RedisMemory) -> None: + content_1 = MemoryContent( + content="I enjoy fruits like apples, oranges, and bananas.", mime_type=MemoryMimeType.TEXT, metadata={} + ) + await semantic_memory.add(content_1) + + # find matches with a similar query + memories = await semantic_memory.query("Fruits that I like.") + assert len(memories.results) == 1 + + # don't return anything for dissimilar queries + no_memories = await semantic_memory.query("The king of England") + assert len(no_memories.results) == 0 + + # match multiple relevant memories + content_2 = MemoryContent( + content="I also like mangos and pineapples.", + mime_type=MemoryMimeType.TEXT, + metadata={"description": "additional info"}, + ) + await semantic_memory.add(content_2) + + memories = await semantic_memory.query("Fruits that I like.") + assert len(memories.results) == 2 + assert memories.results[0].metadata == {} + assert memories.results[1].metadata == {"description": "additional info"} + + +@pytest.mark.asyncio +@pytest.mark.skipif(not redis_available(), reason="Redis instance not available locally") +async def test_clear(semantic_memory: RedisMemory) -> None: + content = MemoryContent(content="I enjoy fruits like apples, oranges, and bananas.", mime_type=MemoryMimeType.TEXT) + await semantic_memory.add(content) + + # find matches with a similar query + memories = await semantic_memory.query("Fruits that I like.") + assert len(memories.results) == 1 + + await semantic_memory.clear() + # don't return anything for dissimilar queries + no_memories = await semantic_memory.query("Fruits that I like.") + assert len(no_memories.results) == 0 + + +@pytest.mark.asyncio +@pytest.mark.skipif(not redis_available(), reason="Redis instance not available locally") +async def test_close(semantic_config: RedisMemoryConfig) -> None: + semantic_memory = RedisMemory(semantic_config) + content = MemoryContent(content="This sentence should be deleted.", mime_type=MemoryMimeType.TEXT) + await semantic_memory.add(content) + + await semantic_memory.close() + + with pytest.raises(RedisSearchError): + _ = await semantic_memory.query("This query should fail.") + + +## INTEGRATION TESTS ## +@pytest.mark.asyncio +@pytest.mark.skipif(not redis_available(), reason="Redis instance not available locally") +async def test_basic_workflow(semantic_config: RedisMemoryConfig) -> None: + """Test basic memory operations with semantic memory.""" + memory = RedisMemory(config=semantic_config) + await memory.clear() + + await memory.add( + MemoryContent( + content="Virginia Tech is the best engineering university in the state.", + mime_type=MemoryMimeType.TEXT, + metadata={"topic": "higher education", "department": "engineering"}, + ) + ) + + results = await memory.query("Which engineering university should I attend?") + assert len(results.results) == 1 + assert any("engineering" in str(r.content) for r in results.results) + assert all(isinstance(r.metadata, dict) for r in results.results if r.metadata) + + await memory.close() + + +@pytest.mark.asyncio +@pytest.mark.skipif(not redis_available(), reason="Redis instance not available locally") +async def test_text_memory_type(semantic_memory: RedisMemory) -> None: + await semantic_memory.clear() + + # Test text content + text_content = MemoryContent(content="Simple text content for testing", mime_type=MemoryMimeType.TEXT) + await semantic_memory.add(text_content) + + # Query for text content + results = await semantic_memory.query("simple text content") + assert len(results.results) > 0 + assert any("Simple text content" in str(r.content) for r in results.results) + + +@pytest.mark.asyncio +@pytest.mark.skipif(not redis_available(), reason="Redis instance not available locally") +async def test_json_memory_type(semantic_memory: RedisMemory) -> None: + await semantic_memory.clear() + + json_data = {"title": "Hitchhiker's Guide to the Galaxy", "The answer to life, the universe and everything.": 42} + await semantic_memory.add( + MemoryContent(content=json_data, mime_type=MemoryMimeType.JSON, metadata={"author": "Douglas Adams"}) + ) + + results = await semantic_memory.query("what is the ultimate question of the universe?") + assert results.results[0].content == json_data + + # meta data should not be searched + results = await semantic_memory.query("who is Douglas Adams?") + assert len(results.results) == 0 + + # test we can't query with JSON also + with pytest.raises(TypeError): + results = await semantic_memory.query({"question": "what is the ultimate question of the universe?"}) # type: ignore[arg-type] + + # but we can if the JSON is within a MemoryContent container + results = await semantic_memory.query( + MemoryContent( + content={"question": "what is the ultimate question of the universe?"}, mime_type=MemoryMimeType.JSON + ) + ) + assert results.results[0].content == json_data + + +@pytest.mark.asyncio +@pytest.mark.skipif(not redis_available(), reason="Redis instance not available locally") +async def test_markdown_memory_type(semantic_memory: RedisMemory) -> None: + await semantic_memory.clear() + + markdown_data = """ + This is an H1 header + ============ + + Paragraphs are separated by a blank line. + + *Italics are within asteriks*, **bold text is within two asterisks**, + while `monospace is within back tics`. + + Itemized lists are made with indented asterisks: + + * this one + * that one + * the next one + + > Block quotes are make with arrows + > like this. + > + > They can span multiple paragraphs, + > if you like. + + Unicode is supported. ☺ + """ + + await semantic_memory.add( + MemoryContent(content=markdown_data, mime_type=MemoryMimeType.MARKDOWN, metadata={"type": "markdown example"}) + ) + + results = await semantic_memory.query("how can I make itemized lists, or italicize text with asterisks?") + assert results.results[0].content == markdown_data + + # test we can query with markdown interpreted as a text string also + results = await semantic_memory.query("") + + # we can also if the markdown is within a MemoryContent container + results = await semantic_memory.query( + MemoryContent( + content="**bold text is within 2 asterisks**, and *italics are within 1 asterisk*", + mime_type=MemoryMimeType.MARKDOWN, + ) + ) + assert results.results[0].content == markdown_data + + +@pytest.mark.asyncio +@pytest.mark.skipif(not redis_available(), reason="Redis instance not available locally") +async def test_query_arguments(semantic_memory: RedisMemory) -> None: + # test that we can utilize the optional query arguments top_k and distance_threshold + await semantic_memory.clear() + + await semantic_memory.add(MemoryContent(content="my favorite fruit are apples", mime_type=MemoryMimeType.TEXT)) + await semantic_memory.add(MemoryContent(content="I also like cherries", mime_type=MemoryMimeType.TEXT)) + await semantic_memory.add(MemoryContent(content="I like plums as well", mime_type=MemoryMimeType.TEXT)) + + # default search + results = await semantic_memory.query("what fruits do I like?") + assert len(results.results) == 3 + + # limit search to 2 results + results = await semantic_memory.query("what fruits do I like?", top_k=2) + assert len(results.results) == 2 + + # limit search to only close matches + results = await semantic_memory.query("my favorite fruit are what?", distance_threshold=0.2) + assert len(results.results) == 1 diff --git a/python/packages/autogen-ext/tests/memory/test_text_canvas_memory.py b/python/packages/autogen-ext/tests/memory/test_text_canvas_memory.py new file mode 100644 index 000000000000..0a7092acf9ed --- /dev/null +++ b/python/packages/autogen-ext/tests/memory/test_text_canvas_memory.py @@ -0,0 +1,121 @@ +import difflib + +import pytest +from autogen_core import CancellationToken +from autogen_core.model_context import UnboundedChatCompletionContext +from autogen_ext.memory.canvas import TextCanvasMemory +from autogen_ext.memory.canvas._canvas_writer import ( + ApplyPatchArgs, + UpdateFileArgs, +) + + +# ── Fixtures ───────────────────────────────────────────────────────────────────── +@pytest.fixture() +def story_v1() -> str: + # Extracted (slightly trimmed) from the sample output + return ( + "# The Bunny and the Sunflower\n\n" + "## Beginning\n" + "Once upon a time, in a bright and cheerful meadow, Bella the bunny came " + "across **a beautiful sunflower** waving in the sunshine.\n" + ) + + +@pytest.fixture() +def story_v2(story_v1: str) -> str: + # A small edit: give the sunflower a name (mirrors the first patch in the log) + return story_v1.replace( + "a beautiful sunflower", + "a beautiful sunflower named Sunny", + ) + + +@pytest.fixture() +def memory() -> TextCanvasMemory: + return TextCanvasMemory() + + +# ── Tests ──────────────────────────────────────────────────────────────────────── +@pytest.mark.asyncio +async def test_canvas_initial_state(memory: TextCanvasMemory) -> None: + assert memory.canvas.list_files() == {} + snapshot = memory.canvas.get_all_contents_for_context() + assert snapshot.startswith("=== CANVAS FILES ===") + + +@pytest.mark.asyncio +async def test_update_file_tool_creates_file( + memory: TextCanvasMemory, + story_v1: str, +) -> None: + update_tool = memory.get_update_file_tool() + + await update_tool.run( + UpdateFileArgs(filename="story.md", new_content=story_v1), + CancellationToken(), + ) + + assert memory.canvas.get_latest_content("story.md") == story_v1 + assert memory.canvas.list_files()["story.md"] == 1 + + +@pytest.mark.asyncio +async def test_apply_patch_increments_revision( + memory: TextCanvasMemory, + story_v1: str, + story_v2: str, +) -> None: + # Set up revision 1 + await memory.get_update_file_tool().run( + UpdateFileArgs(filename="story.md", new_content=story_v1), + CancellationToken(), + ) + + # Create a unified diff for the patch tool + diff_text = "".join( + difflib.unified_diff( + story_v1.splitlines(keepends=True), + story_v2.splitlines(keepends=True), + fromfile="story.md", + tofile="story.md", + ) + ) + + # Apply the patch → revision 2 + await memory.get_apply_patch_tool().run( + ApplyPatchArgs(filename="story.md", patch_text=diff_text), + CancellationToken(), + ) + + assert memory.canvas.get_latest_content("story.md") == story_v2 + # The revision number should now be 2 + assert memory.canvas.list_files()["story.md"] == 2 + # And the diff history should contain exactly one patch + assert len(memory.canvas.get_revision_diffs("story.md")) == 1 + + +@pytest.mark.asyncio +async def test_update_context_injects_snapshot( + memory: TextCanvasMemory, + story_v2: str, +) -> None: + # Seed with some content + await memory.get_update_file_tool().run( + UpdateFileArgs(filename="story.md", new_content=story_v2), + CancellationToken(), + ) + + chat_ctx = UnboundedChatCompletionContext() + result = await memory.update_context(chat_ctx) + + # A single SystemMessage should have been added to the context + assert len(chat_ctx._messages) == 1 # type: ignore + injected_text = chat_ctx._messages[0].content # type: ignore + assert "=== CANVAS FILES ===" in injected_text + assert "story.md" in injected_text + + # The UpdateContextResult should surface the same snapshot via MemoryContent + assert result.memories.results + assert isinstance(result.memories.results[0].content, str) + assert story_v2.strip() in result.memories.results[0].content diff --git a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py index 4f4135b4b0c8..9bc6d4315079 100644 --- a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py +++ b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py @@ -2,6 +2,7 @@ import logging import os from typing import List, Sequence +from unittest.mock import AsyncMock, MagicMock, patch import pytest from autogen_core import CancellationToken, FunctionCall @@ -10,12 +11,17 @@ CreateResult, FunctionExecutionResult, FunctionExecutionResultMessage, + ModelInfo, SystemMessage, UserMessage, ) from autogen_core.models._types import LLMMessage from autogen_core.tools import FunctionTool -from autogen_ext.models.anthropic import AnthropicChatCompletionClient +from autogen_ext.models.anthropic import ( + AnthropicBedrockChatCompletionClient, + AnthropicChatCompletionClient, + BedrockInfo, +) def _pass_function(input: str) -> str: @@ -29,7 +35,153 @@ def _add_numbers(a: int, b: int) -> int: @pytest.mark.asyncio -async def test_anthropic_serialization_api_key() -> None: +async def test_mock_tool_choice_specific_tool() -> None: + """Test tool_choice parameter with a specific tool using mocks.""" + # Create mock client and response + mock_client = AsyncMock() + mock_message = MagicMock() + mock_message.content = [MagicMock(type="tool_use", name="process_text", input={"input": "hello"}, id="call_123")] + mock_message.usage.input_tokens = 10 + mock_message.usage.output_tokens = 5 + + mock_client.messages.create.return_value = mock_message + + # Create real client but patch the underlying Anthropic client + client = AnthropicChatCompletionClient( + model="claude-3-haiku-20240307", + api_key="test-key", + ) + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="process_text") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") + + messages: List[LLMMessage] = [ + UserMessage(content="Process the text 'hello'.", source="user"), + ] + + with patch.object(client, "_client", mock_client): + await client.create( + messages=messages, + tools=[pass_tool, add_tool], + tool_choice=pass_tool, # Force use of specific tool + ) + + # Verify the correct API call was made + mock_client.messages.create.assert_called_once() + call_args = mock_client.messages.create.call_args + + # Check that tool_choice was set correctly + assert "tool_choice" in call_args.kwargs + assert call_args.kwargs["tool_choice"] == {"type": "tool", "name": "process_text"} + + +@pytest.mark.asyncio +async def test_mock_tool_choice_auto() -> None: + """Test tool_choice parameter with 'auto' setting using mocks.""" + # Create mock client and response + mock_client = AsyncMock() + mock_message = MagicMock() + mock_message.content = [MagicMock(type="tool_use", name="add_numbers", input={"a": 1, "b": 2}, id="call_123")] + mock_message.usage.input_tokens = 10 + mock_message.usage.output_tokens = 5 + + mock_client.messages.create.return_value = mock_message + + # Create real client but patch the underlying Anthropic client + client = AnthropicChatCompletionClient( + model="claude-3-haiku-20240307", + api_key="test-key", + ) + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="process_text") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") + + messages: List[LLMMessage] = [ + UserMessage(content="Add 1 and 2.", source="user"), + ] + + with patch.object(client, "_client", mock_client): + await client.create( + messages=messages, + tools=[pass_tool, add_tool], + tool_choice="auto", # Let model choose + ) + + # Verify the correct API call was made + mock_client.messages.create.assert_called_once() + call_args = mock_client.messages.create.call_args + + # Check that tool_choice was set correctly + assert "tool_choice" in call_args.kwargs + assert call_args.kwargs["tool_choice"] == {"type": "auto"} + + +@pytest.mark.asyncio +async def test_mock_tool_choice_none() -> None: + """Test tool_choice parameter when no tools are provided - tool_choice should not be included.""" + # Create mock client and response + mock_client = AsyncMock() + mock_message = MagicMock() + mock_message.content = [MagicMock(type="text", text="I can help you with that.")] + mock_message.usage.input_tokens = 10 + mock_message.usage.output_tokens = 5 + + mock_client.messages.create.return_value = mock_message + + # Create real client but patch the underlying Anthropic client + client = AnthropicChatCompletionClient( + model="claude-3-haiku-20240307", + api_key="test-key", + ) + + messages: List[LLMMessage] = [ + UserMessage(content="Hello there.", source="user"), + ] + + with patch.object(client, "_client", mock_client): + await client.create( + messages=messages, + # No tools provided - tool_choice should not be included in API call + ) + + # Verify the correct API call was made + mock_client.messages.create.assert_called_once() + call_args = mock_client.messages.create.call_args + + # Check that tool_choice was not set when no tools are provided + assert "tool_choice" not in call_args.kwargs + + +@pytest.mark.asyncio +async def test_mock_tool_choice_validation_error() -> None: + """Test tool_choice validation with invalid tool reference.""" + client = AnthropicChatCompletionClient( + model="claude-3-haiku-20240307", + api_key="test-key", + ) + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="process_text") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") + different_tool = FunctionTool(_pass_function, description="Different tool", name="different_tool") + + messages: List[LLMMessage] = [ + UserMessage(content="Hello there.", source="user"), + ] + + # Test with a tool that's not in the tools list + with pytest.raises(ValueError, match="tool_choice references 'different_tool' but it's not in the available tools"): + await client.create( + messages=messages, + tools=[pass_tool, add_tool], + tool_choice=different_tool, # This tool is not in the tools list + ) + + +@pytest.mark.asyncio +async def test_mock_serialization_api_key() -> None: client = AnthropicChatCompletionClient( model="claude-3-haiku-20240307", # Use haiku for faster/cheaper testing api_key="sk-password", @@ -46,6 +198,29 @@ async def test_anthropic_serialization_api_key() -> None: client2 = AnthropicChatCompletionClient.load_component(config) assert client2 + bedrock_client = AnthropicBedrockChatCompletionClient( + model="claude-3-haiku-20240307", # Use haiku for faster/cheaper testing + api_key="sk-password", + model_info=ModelInfo( + vision=False, function_calling=True, json_output=False, family="unknown", structured_output=True + ), + bedrock_info=BedrockInfo( + aws_access_key="", + aws_secret_key="", + aws_session_token="", + aws_region="", + ), + ) + assert bedrock_client + bedrock_config = bedrock_client.dump_component() + assert bedrock_config + assert "sk-password" not in str(bedrock_config) + serialized_bedrock_config = bedrock_config.model_dump_json() + assert serialized_bedrock_config + assert "sk-password" not in serialized_bedrock_config + bedrock_client2 = AnthropicBedrockChatCompletionClient.load_component(bedrock_config) + assert bedrock_client2 + @pytest.mark.asyncio async def test_anthropic_basic_completion(caplog: pytest.LogCaptureFixture) -> None: @@ -314,16 +489,12 @@ async def test_anthropic_multimodal() -> None: @pytest.mark.asyncio -async def test_anthropic_serialization() -> None: +async def test_mock_serialization() -> None: """Test serialization and deserialization of component.""" - api_key = os.getenv("ANTHROPIC_API_KEY") - if not api_key: - pytest.skip("ANTHROPIC_API_KEY not found in environment variables") - client = AnthropicChatCompletionClient( model="claude-3-haiku-20240307", - api_key=api_key, + api_key="api-key", ) # Serialize and deserialize @@ -334,3 +505,497 @@ async def test_anthropic_serialization() -> None: loaded_model_client = AnthropicChatCompletionClient.load_component(model_client_config) assert loaded_model_client is not None assert isinstance(loaded_model_client, AnthropicChatCompletionClient) + + +@pytest.mark.asyncio +async def test_anthropic_message_serialization_with_tools(caplog: pytest.LogCaptureFixture) -> None: + """Test that complex messages with tool calls are properly serialized in logs.""" + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not found in environment variables") + + # Use existing tools from the test file + pass_tool = FunctionTool(_pass_function, description="Process input text", name="process_text") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") + + client = AnthropicChatCompletionClient( + model="claude-3-haiku-20240307", + api_key=api_key, + ) + + # Set up logging capture - capture all loggers + with caplog.at_level(logging.INFO): + # Make a request that should trigger a tool call + await client.create( + messages=[ + SystemMessage(content="Use the tools available to help the user."), + UserMessage(content="Process the text 'hello world' using the process_text tool.", source="user"), + ], + tools=[pass_tool, add_tool], + ) + + # Look for any log containing serialized messages, not just with 'LLMCallEvent' + serialized_message_logs = [ + record for record in caplog.records if '"messages":' in str(record.msg) or "messages" in str(record.msg) + ] + + # Verify we have at least one log with serialized messages + assert len(serialized_message_logs) > 0, "No logs with serialized messages found" + + +@pytest.mark.asyncio +async def test_anthropic_muliple_system_message() -> None: + """Test multiple system messages in a single request.""" + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not found in environment variables") + + client = AnthropicChatCompletionClient( + model="claude-3-haiku-20240307", + api_key=api_key, + ) + # Test multiple system messages + messages: List[LLMMessage] = [ + SystemMessage(content="When you say anything Start with 'FOO'"), + SystemMessage(content="When you say anything End with 'BAR'"), + UserMessage(content="Just say '.'", source="user"), + ] + + result = await client.create(messages=messages) + result_content = result.content + assert isinstance(result_content, str) + result_content = result_content.strip() + assert result_content[:3] == "FOO" + assert result_content[-3:] == "BAR" + + +def test_mock_merge_continuous_system_messages() -> None: + """Tests merging of continuous system messages.""" + client = AnthropicChatCompletionClient(model="claude-3-haiku-20240307", api_key="fake-api-key") + + messages: List[LLMMessage] = [ + SystemMessage(content="System instruction 1"), + SystemMessage(content="System instruction 2"), + UserMessage(content="User question", source="user"), + ] + + merged_messages = client._merge_system_messages(messages) # pyright: ignore[reportPrivateUsage] + # The method is protected, but we need to test it + + # 병합 후 2개 메시지만 남아야 함 (시스템 1개, 사용자 1개) + assert len(merged_messages) == 2 + + # 첫 번째 메시지는 병합된 시스템 메시지여야 함 + assert isinstance(merged_messages[0], SystemMessage) + assert merged_messages[0].content == "System instruction 1\nSystem instruction 2" + + # 두 번째 메시지는 사용자 메시지여야 함 + assert isinstance(merged_messages[1], UserMessage) + assert merged_messages[1].content == "User question" + + +def test_mock_merge_single_system_message() -> None: + """Tests that a single system message remains unchanged.""" + client = AnthropicChatCompletionClient(model="claude-3-haiku-20240307", api_key="fake-api-key") + + messages: List[LLMMessage] = [ + SystemMessage(content="Single system instruction"), + UserMessage(content="User question", source="user"), + ] + + merged_messages = client._merge_system_messages(messages) # pyright: ignore[reportPrivateUsage] + # The method is protected, but we need to test it + + # 메시지 개수는 변하지 않아야 함 + assert len(merged_messages) == 2 + + # 시스템 메시지 내용은 변하지 않아야 함 + assert isinstance(merged_messages[0], SystemMessage) + assert merged_messages[0].content == "Single system instruction" + + +def test_mock_merge_no_system_messages() -> None: + """Tests behavior when there are no system messages.""" + client = AnthropicChatCompletionClient(model="claude-3-haiku-20240307", api_key="fake-api-key") + + messages: List[LLMMessage] = [ + UserMessage(content="User question without system", source="user"), + ] + + merged_messages = client._merge_system_messages(messages) # pyright: ignore[reportPrivateUsage] + # The method is protected, but we need to test it + + # 메시지 개수는 변하지 않아야 함 + assert len(merged_messages) == 1 + + # 유일한 메시지는 사용자 메시지여야 함 + assert isinstance(merged_messages[0], UserMessage) + assert merged_messages[0].content == "User question without system" + + +def test_mock_merge_non_continuous_system_messages() -> None: + """Tests that an error is raised for non-continuous system messages.""" + client = AnthropicChatCompletionClient(model="claude-3-haiku-20240307", api_key="fake-api-key") + + messages: List[LLMMessage] = [ + SystemMessage(content="First group 1"), + SystemMessage(content="First group 2"), + UserMessage(content="Middle user message", source="user"), + SystemMessage(content="Second group 1"), + SystemMessage(content="Second group 2"), + ] + + # 연속적이지 않은 시스템 메시지는 에러를 발생시켜야 함 + with pytest.raises(ValueError, match="Multiple and Not continuous system messages are not supported"): + client._merge_system_messages(messages) # pyright: ignore[reportPrivateUsage] + # The method is protected, but we need to test it + + +def test_mock_merge_system_messages_empty() -> None: + """Tests that empty message list is handled properly.""" + client = AnthropicChatCompletionClient(model="claude-3-haiku-20240307", api_key="fake-api-key") + + merged_messages = client._merge_system_messages([]) # pyright: ignore[reportPrivateUsage] + # The method is protected, but we need to test it + assert len(merged_messages) == 0 + + +def test_mock_merge_system_messages_with_special_characters() -> None: + """Tests system message merging with special characters and formatting.""" + client = AnthropicChatCompletionClient(model="claude-3-haiku-20240307", api_key="fake-api-key") + + messages: List[LLMMessage] = [ + SystemMessage(content="Line 1\nWith newline"), + SystemMessage(content="Line 2 with *formatting*"), + SystemMessage(content="Line 3 with `code`"), + UserMessage(content="Question", source="user"), + ] + + merged_messages = client._merge_system_messages(messages) # pyright: ignore[reportPrivateUsage] + # The method is protected, but we need to test it + assert len(merged_messages) == 2 + + system_message = merged_messages[0] + assert isinstance(system_message, SystemMessage) + assert system_message.content == "Line 1\nWith newline\nLine 2 with *formatting*\nLine 3 with `code`" + + +def test_mock_merge_system_messages_with_whitespace() -> None: + """Tests system message merging with extra whitespace.""" + client = AnthropicChatCompletionClient(model="claude-3-haiku-20240307", api_key="fake-api-key") + + messages: List[LLMMessage] = [ + SystemMessage(content=" Message with leading spaces "), + SystemMessage(content="\nMessage with leading newline\n"), + UserMessage(content="Question", source="user"), + ] + + merged_messages = client._merge_system_messages(messages) # pyright: ignore[reportPrivateUsage] + # The method is protected, but we need to test it + assert len(merged_messages) == 2 + + system_message = merged_messages[0] + assert isinstance(system_message, SystemMessage) + # strip()은 내부에서 발생하지 않지만 최종 결과에서는 줄바꿈이 유지됨 + assert system_message.content == " Message with leading spaces \n\nMessage with leading newline" + + +def test_mock_merge_system_messages_message_order() -> None: + """Tests that message order is preserved after merging.""" + client = AnthropicChatCompletionClient(model="claude-3-haiku-20240307", api_key="fake-api-key") + + messages: List[LLMMessage] = [ + UserMessage(content="Question 1", source="user"), + SystemMessage(content="Instruction 1"), + SystemMessage(content="Instruction 2"), + UserMessage(content="Question 2", source="user"), + AssistantMessage(content="Answer", source="assistant"), + ] + + merged_messages = client._merge_system_messages(messages) # pyright: ignore[reportPrivateUsage] + # The method is protected, but we need to test it + assert len(merged_messages) == 4 + + # 첫 번째 메시지는 UserMessage여야 함 + assert isinstance(merged_messages[0], UserMessage) + assert merged_messages[0].content == "Question 1" + + # 두 번째 메시지는 병합된 SystemMessage여야 함 + assert isinstance(merged_messages[1], SystemMessage) + assert merged_messages[1].content == "Instruction 1\nInstruction 2" + + # 나머지 메시지는 순서대로 유지되어야 함 + assert isinstance(merged_messages[2], UserMessage) + assert merged_messages[2].content == "Question 2" + assert isinstance(merged_messages[3], AssistantMessage) + assert merged_messages[3].content == "Answer" + + +def test_mock_merge_system_messages_multiple_groups() -> None: + """Tests that multiple separate groups of system messages raise an error.""" + client = AnthropicChatCompletionClient(model="claude-3-haiku-20240307", api_key="fake-api-key") + + # 연속되지 않은 시스템 메시지: 사용자 메시지로 분리된 두 그룹 + messages: List[LLMMessage] = [ + SystemMessage(content="Group 1 - message 1"), + UserMessage(content="Interrupting user message", source="user"), + SystemMessage(content="Group 2 - message 1"), + ] + + with pytest.raises(ValueError, match="Multiple and Not continuous system messages are not supported"): + client._merge_system_messages(messages) # pyright: ignore[reportPrivateUsage] + # The method is protected, but we need to test it + + +def test_mock_merge_system_messages_no_duplicates() -> None: + """Tests that identical system messages are still merged properly.""" + client = AnthropicChatCompletionClient(model="claude-3-haiku-20240307", api_key="fake-api-key") + + messages: List[LLMMessage] = [ + SystemMessage(content="Same instruction"), + SystemMessage(content="Same instruction"), # 중복된 내용 + UserMessage(content="Question", source="user"), + ] + + merged_messages = client._merge_system_messages(messages) # pyright: ignore[reportPrivateUsage] + # The method is protected, but we need to test it + assert len(merged_messages) == 2 + + # 첫 번째 메시지는 병합된 시스템 메시지여야 함 + assert isinstance(merged_messages[0], SystemMessage) + # 중복된 내용도 그대로 병합됨 + assert merged_messages[0].content == "Same instruction\nSame instruction" + + +@pytest.mark.asyncio +async def test_anthropic_empty_assistant_content_string() -> None: + """Test that an empty assistant content string is handled correctly.""" + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not found in environment variables") + + client = AnthropicChatCompletionClient( + model="claude-3-haiku-20240307", + api_key=api_key, + ) + + # Test empty assistant content string + result = await client.create( + messages=[ + UserMessage(content="Say something", source="user"), + AssistantMessage(content="", source="assistant"), + ] + ) + + # Verify we got a response + assert isinstance(result.content, str) + assert len(result.content) > 0 + + +@pytest.mark.asyncio +async def test_anthropic_trailing_whitespace_at_last_assistant_content() -> None: + """Test that an empty assistant content string is handled correctly.""" + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not found in environment variables") + + client = AnthropicChatCompletionClient( + model="claude-3-haiku-20240307", + api_key=api_key, + ) + + messages: list[LLMMessage] = [ + UserMessage(content="foo", source="user"), + UserMessage(content="bar", source="user"), + AssistantMessage(content="foobar ", source="assistant"), + ] + + result = await client.create(messages=messages) + assert isinstance(result.content, str) + + +def test_mock_rstrip_trailing_whitespace_at_last_assistant_content() -> None: + messages: list[LLMMessage] = [ + UserMessage(content="foo", source="user"), + UserMessage(content="bar", source="user"), + AssistantMessage(content="foobar ", source="assistant"), + ] + + # This will crash if _rstrip_railing_whitespace_at_last_assistant_content is not applied to "content" + dummy_client = AnthropicChatCompletionClient(model="claude-3-5-haiku-20241022", api_key="dummy-key") + result = dummy_client._rstrip_last_assistant_message(messages) # pyright: ignore[reportPrivateUsage] + + assert isinstance(result[-1].content, str) + assert result[-1].content == "foobar" + + +@pytest.mark.asyncio +async def test_anthropic_tool_choice_with_actual_api() -> None: + """Test tool_choice parameter with actual Anthropic API endpoints.""" + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not found in environment variables") + + client = AnthropicChatCompletionClient( + model="claude-3-haiku-20240307", + api_key=api_key, + ) + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="process_text") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") + + # Test 1: tool_choice with specific tool + messages: List[LLMMessage] = [ + SystemMessage(content="Use the tools as needed to help the user."), + UserMessage(content="Process the text 'hello world' using the process_text tool.", source="user"), + ] + + result = await client.create( + messages=messages, + tools=[pass_tool, add_tool], + tool_choice=pass_tool, # Force use of specific tool + ) + + # Verify we got a tool call for the specified tool + assert isinstance(result.content, list) + assert len(result.content) >= 1 + assert isinstance(result.content[0], FunctionCall) + assert result.content[0].name == "process_text" + + # Test 2: tool_choice="auto" with tools + auto_messages: List[LLMMessage] = [ + SystemMessage(content="Use the tools as needed to help the user."), + UserMessage(content="Add the numbers 5 and 3.", source="user"), + ] + + auto_result = await client.create( + messages=auto_messages, + tools=[pass_tool, add_tool], + tool_choice="auto", # Let model choose + ) + + # Should get a tool call, likely for add_numbers + assert isinstance(auto_result.content, list) + assert len(auto_result.content) >= 1 + assert isinstance(auto_result.content[0], FunctionCall) + # Model should choose add_numbers for addition task + assert auto_result.content[0].name == "add_numbers" + + # Test 3: No tools provided - should not include tool_choice in API call + no_tools_messages: List[LLMMessage] = [ + UserMessage(content="What is the capital of France?", source="user"), + ] + + no_tools_result = await client.create(messages=no_tools_messages) + + # Should get a text response without tool calls + assert isinstance(no_tools_result.content, str) + assert "paris" in no_tools_result.content.lower() + + # Test 4: tool_choice="required" with tools + required_messages: List[LLMMessage] = [ + SystemMessage(content="You must use one of the available tools to help the user."), + UserMessage(content="Help me with something.", source="user"), + ] + + required_result = await client.create( + messages=required_messages, + tools=[pass_tool, add_tool], + tool_choice="required", # Force tool usage + ) + + # Should get a tool call (model forced to use a tool) + assert isinstance(required_result.content, list) + assert len(required_result.content) >= 1 + assert isinstance(required_result.content[0], FunctionCall) + + +@pytest.mark.asyncio +async def test_anthropic_tool_choice_streaming_with_actual_api() -> None: + """Test tool_choice parameter with streaming using actual Anthropic API endpoints.""" + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not found in environment variables") + + client = AnthropicChatCompletionClient( + model="claude-3-haiku-20240307", + api_key=api_key, + ) + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="process_text") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") + + # Test streaming with tool_choice + messages: List[LLMMessage] = [ + SystemMessage(content="Use the tools as needed to help the user."), + UserMessage(content="Process the text 'streaming test' using the process_text tool.", source="user"), + ] + + chunks: List[str | CreateResult] = [] + async for chunk in client.create_stream( + messages=messages, + tools=[pass_tool, add_tool], + tool_choice=pass_tool, # Force use of specific tool + ): + chunks.append(chunk) + + # Verify we got chunks and a final result + assert len(chunks) > 0 + final_result = chunks[-1] + assert isinstance(final_result, CreateResult) + + # Should get a tool call for the specified tool + assert isinstance(final_result.content, list) + assert len(final_result.content) >= 1 + assert isinstance(final_result.content[0], FunctionCall) + assert final_result.content[0].name == "process_text" + + # Test streaming without tools - should not include tool_choice + no_tools_messages: List[LLMMessage] = [ + UserMessage(content="Tell me a short fact about cats.", source="user"), + ] + + no_tools_chunks: List[str | CreateResult] = [] + async for chunk in client.create_stream(messages=no_tools_messages): + no_tools_chunks.append(chunk) + + # Should get text response + assert len(no_tools_chunks) > 0 + final_no_tools_result = no_tools_chunks[-1] + assert isinstance(final_no_tools_result, CreateResult) + assert isinstance(final_no_tools_result.content, str) + assert len(final_no_tools_result.content) > 0 + + +@pytest.mark.asyncio +async def test_anthropic_tool_choice_none_value_with_actual_api() -> None: + """Test tool_choice="none" with actual Anthropic API endpoints.""" + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not found in environment variables") + + client = AnthropicChatCompletionClient( + model="claude-3-haiku-20240307", + api_key=api_key, + ) + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="process_text") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") + + # Test tool_choice="none" - should not use tools even when available + messages: List[LLMMessage] = [ + SystemMessage(content="Answer the user's question directly without using tools."), + UserMessage(content="What is 2 + 2?", source="user"), + ] + + result = await client.create( + messages=messages, + tools=[pass_tool, add_tool], + tool_choice="none", # Disable tool usage + ) + + # Should get a text response, not tool calls + assert isinstance(result.content, str) diff --git a/python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py b/python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py index 2f4f02aeaf13..dfc7af07302f 100644 --- a/python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py +++ b/python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py @@ -3,11 +3,12 @@ import os from datetime import datetime from typing import Any, AsyncGenerator, List, Type, Union -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest from autogen_core import CancellationToken, FunctionCall, Image from autogen_core.models import CreateResult, ModelFamily, UserMessage +from autogen_core.tools import FunctionTool from autogen_ext.models.azure import AzureAIChatCompletionClient from autogen_ext.models.azure.config import GITHUB_MODELS_ENDPOINT from azure.ai.inference.aio import ( @@ -570,7 +571,7 @@ async def _mock_thought_with_tool_call_stream( ) mock_client = MagicMock() - mock_client.close = MagicMock() + mock_client.close = AsyncMock() async def mock_complete(*args: Any, **kwargs: Any) -> Any: if kwargs.get("stream", False): @@ -623,3 +624,352 @@ async def test_thought_field_with_tool_calls_streaming( assert final_result.content[0].arguments == '{"foo": "bar"}' assert final_result.thought == "Let me think about what function to call." + + +def _pass_function(input: str) -> str: + """Simple passthrough function.""" + return f"Processed: {input}" + + +def _add_numbers(a: int, b: int) -> int: + """Add two numbers together.""" + return a + b + + +@pytest.fixture +def tool_choice_client(monkeypatch: pytest.MonkeyPatch) -> AzureAIChatCompletionClient: + """ + Returns a client that supports function calling for tool choice tests. + """ + + async def _mock_tool_choice_stream( + *args: Any, **kwargs: Any + ) -> AsyncGenerator[StreamingChatCompletionsUpdate, None]: + mock_chunks_content = ["Hello", " Another Hello", " Yet Another Hello"] + + mock_chunks = [ + StreamingChatChoiceUpdate( + index=0, + finish_reason="stop", + delta=StreamingChatResponseMessageUpdate(role="assistant", content=chunk_content), + ) + for chunk_content in mock_chunks_content + ] + + for mock_chunk in mock_chunks: + await asyncio.sleep(0.01) + yield StreamingChatCompletionsUpdate( + id="id", + choices=[mock_chunk], + created=datetime.now(), + model="model", + usage=CompletionsUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + + mock_client = MagicMock() + mock_client.close = AsyncMock() + + async def mock_complete(*args: Any, **kwargs: Any) -> Any: + stream = kwargs.get("stream", False) + + if not stream: + await asyncio.sleep(0.01) + return ChatCompletions( + id="id", + created=datetime.now(), + model="model", + choices=[ + ChatChoice( + index=0, + finish_reason=CompletionsFinishReason.TOOL_CALLS, + message=ChatResponseMessage( + role="assistant", + content="", + tool_calls=[ + ChatCompletionsToolCall( + id="call_123", + function=AzureFunctionCall(name="process_text", arguments='{"input": "hello"}'), + ) + ], + ), + ) + ], + usage=CompletionsUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + else: + return _mock_tool_choice_stream(*args, **kwargs) + + mock_client.complete = mock_complete + + def mock_new(cls: Type[ChatCompletionsClient], *args: Any, **kwargs: Any) -> MagicMock: + return mock_client + + monkeypatch.setattr(ChatCompletionsClient, "__new__", mock_new) + + return AzureAIChatCompletionClient( + endpoint="endpoint", + credential=AzureKeyCredential("api_key"), + model_info={ + "json_output": False, + "function_calling": True, + "vision": False, + "family": "test", + "structured_output": False, + }, + model="model", + ) + + +@pytest.mark.asyncio +async def test_azure_ai_tool_choice_specific_tool(tool_choice_client: AzureAIChatCompletionClient) -> None: + """Test tool_choice parameter with a specific tool using mocks.""" + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="process_text") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") + + messages = [ + UserMessage(content="Process the text 'hello'.", source="user"), + ] + + result = await tool_choice_client.create( + messages=messages, + tools=[pass_tool, add_tool], + tool_choice=pass_tool, # Force use of specific tool + ) + + # Verify the result + assert result.finish_reason == "function_calls" + assert isinstance(result.content, list) + assert len(result.content) == 1 + assert isinstance(result.content[0], FunctionCall) + assert result.content[0].name == "process_text" + assert result.content[0].arguments == '{"input": "hello"}' + + +@pytest.mark.asyncio +async def test_azure_ai_tool_choice_auto(tool_choice_client: AzureAIChatCompletionClient) -> None: + """Test tool_choice parameter with 'auto' setting using mocks.""" + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="process_text") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") + + messages = [ + UserMessage(content="Add 1 and 2.", source="user"), + ] + + result = await tool_choice_client.create( + messages=messages, + tools=[pass_tool, add_tool], + tool_choice="auto", # Let the model choose + ) + + # Verify the result + assert result.finish_reason == "function_calls" + assert isinstance(result.content, list) + assert len(result.content) == 1 + assert isinstance(result.content[0], FunctionCall) + assert result.content[0].name == "process_text" # Our mock always returns process_text + assert result.content[0].arguments == '{"input": "hello"}' + + +@pytest.fixture +def tool_choice_none_client(monkeypatch: pytest.MonkeyPatch) -> AzureAIChatCompletionClient: + """ + Returns a client that simulates no tool calls for tool_choice='none' tests. + """ + + mock_client = MagicMock() + mock_client.close = AsyncMock() + + async def mock_complete(*args: Any, **kwargs: Any) -> ChatCompletions: + await asyncio.sleep(0.01) + return ChatCompletions( + id="id", + created=datetime.now(), + model="model", + choices=[ + ChatChoice( + index=0, + finish_reason="stop", + message=ChatResponseMessage(role="assistant", content="I can help you with that."), + ) + ], + usage=CompletionsUsage(prompt_tokens=8, completion_tokens=6, total_tokens=14), + ) + + mock_client.complete = mock_complete + + def mock_new(cls: Type[ChatCompletionsClient], *args: Any, **kwargs: Any) -> MagicMock: + return mock_client + + monkeypatch.setattr(ChatCompletionsClient, "__new__", mock_new) + + return AzureAIChatCompletionClient( + endpoint="endpoint", + credential=AzureKeyCredential("api_key"), + model_info={ + "json_output": False, + "function_calling": True, + "vision": False, + "family": "test", + "structured_output": False, + }, + model="model", + ) + + +@pytest.mark.asyncio +async def test_azure_ai_tool_choice_none(tool_choice_none_client: AzureAIChatCompletionClient) -> None: + """Test tool_choice parameter with 'none' setting using mocks.""" + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="process_text") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") + + messages = [ + UserMessage(content="Just say hello.", source="user"), + ] + + result = await tool_choice_none_client.create( + messages=messages, + tools=[pass_tool, add_tool], + tool_choice="none", # Prevent tool usage + ) + + # Verify the result + assert result.finish_reason == "stop" + assert isinstance(result.content, str) + assert result.content == "I can help you with that." + + +@pytest.mark.asyncio +async def test_azure_ai_tool_choice_required(tool_choice_client: AzureAIChatCompletionClient) -> None: + """Test tool_choice parameter with 'required' setting using mocks.""" + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="process_text") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") + + messages = [ + UserMessage(content="Process some text.", source="user"), + ] + + result = await tool_choice_client.create( + messages=messages, + tools=[pass_tool, add_tool], + tool_choice="required", # Force tool usage + ) + + # Verify the result + assert result.finish_reason == "function_calls" + assert isinstance(result.content, list) + assert len(result.content) == 1 + assert isinstance(result.content[0], FunctionCall) + assert result.content[0].name == "process_text" + assert result.content[0].arguments == '{"input": "hello"}' + + +@pytest.fixture +def tool_choice_stream_client(monkeypatch: pytest.MonkeyPatch) -> AzureAIChatCompletionClient: + """ + Returns a client that supports function calling for streaming tool choice tests. + """ + + # Mock tool call for streaming + mock_tool_call = MagicMock() + mock_tool_call.id = "call_123" + mock_tool_call.function = MagicMock() + mock_tool_call.function.name = "process_text" + mock_tool_call.function.arguments = '{"input": "hello"}' + + # First choice with content + first_choice = MagicMock() + first_choice.delta = MagicMock() + first_choice.delta.content = "Let me process this for you." + first_choice.finish_reason = None + + # Tool call choice + tool_call_choice = MagicMock() + tool_call_choice.delta = MagicMock() + tool_call_choice.delta.content = None + tool_call_choice.delta.tool_calls = [mock_tool_call] + tool_call_choice.finish_reason = "function_calls" + + async def _mock_tool_choice_stream( + *args: Any, **kwargs: Any + ) -> AsyncGenerator[StreamingChatCompletionsUpdate, None]: + yield StreamingChatCompletionsUpdate( + id="id", + choices=[first_choice], + created=datetime.now(), + model="model", + ) + + await asyncio.sleep(0.01) + + yield StreamingChatCompletionsUpdate( + id="id", + choices=[tool_call_choice], + created=datetime.now(), + model="model", + usage=CompletionsUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + + mock_client = MagicMock() + mock_client.close = AsyncMock() + + async def mock_complete(*args: Any, **kwargs: Any) -> Any: + if kwargs.get("stream", False): + return _mock_tool_choice_stream(*args, **kwargs) + return None + + mock_client.complete = mock_complete + + def mock_new(cls: Type[ChatCompletionsClient], *args: Any, **kwargs: Any) -> MagicMock: + return mock_client + + monkeypatch.setattr(ChatCompletionsClient, "__new__", mock_new) + + return AzureAIChatCompletionClient( + endpoint="endpoint", + credential=AzureKeyCredential("api_key"), + model_info={ + "json_output": False, + "function_calling": True, + "vision": False, + "family": "test", + "structured_output": False, + }, + model="model", + ) + + +@pytest.mark.asyncio +async def test_azure_ai_tool_choice_specific_tool_streaming( + tool_choice_stream_client: AzureAIChatCompletionClient, +) -> None: + """Test tool_choice parameter with streaming and a specific tool using mocks.""" + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="process_text") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") + + messages = [ + UserMessage(content="Process the text 'hello'.", source="user"), + ] + + chunks: List[Union[str, CreateResult]] = [] + async for chunk in tool_choice_stream_client.create_stream( + messages=messages, + tools=[pass_tool, add_tool], + tool_choice=pass_tool, # Force use of specific tool + ): + chunks.append(chunk) + + # Verify that we got some result + final_result = chunks[-1] + assert isinstance(final_result, CreateResult) + assert final_result.finish_reason == "function_calls" + assert isinstance(final_result.content, list) + assert len(final_result.content) == 1 + assert isinstance(final_result.content[0], FunctionCall) + assert final_result.content[0].name == "process_text" + assert final_result.content[0].arguments == '{"input": "hello"}' + assert final_result.thought == "Let me process this for you." diff --git a/python/packages/autogen-ext/tests/models/test_ollama_chat_completion_client.py b/python/packages/autogen-ext/tests/models/test_ollama_chat_completion_client.py index dec279274eb3..ff88b4a39db5 100644 --- a/python/packages/autogen-ext/tests/models/test_ollama_chat_completion_client.py +++ b/python/packages/autogen-ext/tests/models/test_ollama_chat_completion_client.py @@ -1,6 +1,6 @@ import json import logging -from typing import Any, AsyncGenerator, List, Mapping +from typing import Any, AsyncGenerator, Dict, List, Mapping, Optional import httpx import pytest @@ -13,11 +13,11 @@ FunctionExecutionResultMessage, UserMessage, ) -from autogen_core.tools import FunctionTool +from autogen_core.tools import FunctionTool, ToolSchema from autogen_ext.models.ollama import OllamaChatCompletionClient -from autogen_ext.models.ollama._ollama_client import OLLAMA_VALID_CREATE_KWARGS_KEYS +from autogen_ext.models.ollama._ollama_client import OLLAMA_VALID_CREATE_KWARGS_KEYS, convert_tools from httpx import Response -from ollama import AsyncClient, ChatResponse, Message +from ollama import AsyncClient, ChatResponse, Message, Tool from pydantic import BaseModel @@ -206,6 +206,117 @@ async def _mock_chat(*args: Any, **kwargs: Any) -> ChatResponse: assert create_result.usage.completion_tokens == 12 +@pytest.mark.asyncio +async def test_convert_tools() -> None: + def add(x: int, y: Optional[int]) -> str: + if y is None: + return str(x) + return str(x + y) + + add_tool = FunctionTool(add, description="Add two numbers") + + tool_schema_noparam: ToolSchema = { + "name": "manual_tool", + "description": "A tool defined manually", + "parameters": { + "type": "object", + "properties": { + "param_with_type": {"type": "integer", "description": "An integer param"}, + "param_without_type": {"description": "A param without explicit type"}, + }, + "required": ["param_with_type"], + }, + } + + converted_tools = convert_tools([add_tool, tool_schema_noparam]) + assert len(converted_tools) == 2 + assert isinstance(converted_tools[0].function, Tool.Function) + assert isinstance(converted_tools[0].function.parameters, Tool.Function.Parameters) + assert converted_tools[0].function.parameters.properties is not None + assert converted_tools[0].function.name == add_tool.name + assert converted_tools[0].function.parameters.properties["y"].type == "integer" + + # test it defaults to string + assert isinstance(converted_tools[1].function, Tool.Function) + assert isinstance(converted_tools[1].function.parameters, Tool.Function.Parameters) + assert converted_tools[1].function.parameters.properties is not None + assert converted_tools[1].function.name == "manual_tool" + assert converted_tools[1].function.parameters.properties["param_with_type"].type == "integer" + assert converted_tools[1].function.parameters.properties["param_without_type"].type == "string" + assert converted_tools[1].function.parameters.required == ["param_with_type"] + + +@pytest.mark.asyncio +async def test_create_stream_tools(monkeypatch: pytest.MonkeyPatch) -> None: + def add(x: int, y: int) -> str: + return str(x + y) + + add_tool = FunctionTool(add, description="Add two numbers") + model = "llama3.2" + content_raw = "Hello world! This is a test response. Test response." + + async def _mock_chat(*args: Any, **kwargs: Any) -> AsyncGenerator[ChatResponse, None]: + assert "stream" in kwargs + assert kwargs["stream"] is True + + async def _mock_stream() -> AsyncGenerator[ChatResponse, None]: + chunks = [content_raw[i : i + 5] for i in range(0, len(content_raw), 5)] + # Simulate streaming by yielding chunks of the response + for chunk in chunks[:-1]: + yield ChatResponse( + model=model, + done=False, + message=Message( + role="assistant", + content=chunk, + ), + ) + yield ChatResponse( + model=model, + done=True, + done_reason="stop", + message=Message( + content=chunks[-1], + role="assistant", + tool_calls=[ + Message.ToolCall( + function=Message.ToolCall.Function( + name=add_tool.name, + arguments={"x": 2, "y": 2}, + ), + ), + ], + ), + prompt_eval_count=10, + eval_count=12, + ) + + return _mock_stream() + + monkeypatch.setattr(AsyncClient, "chat", _mock_chat) + client = OllamaChatCompletionClient(model=model) + stream = client.create_stream( + messages=[ + UserMessage(content="hi", source="user"), + ], + tools=[add_tool], + ) + chunks: List[str | CreateResult] = [] + async for chunk in stream: + chunks.append(chunk) + assert len(chunks) > 0 + assert isinstance(chunks[-1], CreateResult) + assert isinstance(chunks[-1].content, list) + assert len(chunks[-1].content) > 0 + assert isinstance(chunks[-1].content[0], FunctionCall) + assert chunks[-1].content[0].name == add_tool.name + assert chunks[-1].content[0].arguments == json.dumps({"x": 2, "y": 2}) + assert chunks[-1].finish_reason == "stop" + assert chunks[-1].usage is not None + assert chunks[-1].usage.prompt_tokens == 10 + assert chunks[-1].usage.completion_tokens == 12 + + @pytest.mark.asyncio async def test_create_structured_output(monkeypatch: pytest.MonkeyPatch) -> None: class ResponseType(BaseModel): @@ -459,7 +570,7 @@ class ResponseType(BaseModel): @pytest.mark.asyncio -@pytest.mark.parametrize("model", ["qwen2.5:0.5b", "llama3.2:1b"]) +@pytest.mark.parametrize("model", ["qwen2.5:0.5b", "llama3.2:1b", "qwen3:0.6b"]) async def test_ollama_create_tools(model: str, ollama_client: OllamaChatCompletionClient) -> None: def add(x: int, y: int) -> str: return str(x + y) @@ -479,6 +590,7 @@ def add(x: int, y: int) -> str: assert len(create_result.content) > 0 assert isinstance(create_result.content[0], FunctionCall) assert create_result.content[0].name == add_tool.name + assert create_result.content[0].arguments == json.dumps({"x": 2, "y": 2}) assert create_result.finish_reason == "function_calls" execution_result = FunctionExecutionResult( @@ -541,9 +653,8 @@ def add(x: int, y: int) -> str: assert ResponseType.model_validate_json(create_result.thought) -@pytest.mark.skip("TODO: Fix streaming with tools") @pytest.mark.asyncio -@pytest.mark.parametrize("model", ["qwen2.5:0.5b", "llama3.2:1b"]) +@pytest.mark.parametrize("model", ["qwen2.5:0.5b", "llama3.2:1b", "qwen3:0.6b"]) async def test_ollama_create_stream_tools(model: str, ollama_client: OllamaChatCompletionClient) -> None: def add(x: int, y: int) -> str: return str(x + y) @@ -569,35 +680,636 @@ def add(x: int, y: int) -> str: assert len(create_result.content) > 0 assert isinstance(create_result.content[0], FunctionCall) assert create_result.content[0].name == add_tool.name - assert create_result.finish_reason == "function_calls" + assert create_result.content[0].arguments == json.dumps({"x": 2, "y": 2}) + assert create_result.finish_reason == "stop" + assert create_result.usage is not None + assert create_result.usage.prompt_tokens == 10 + assert create_result.usage.completion_tokens == 12 - execution_result = FunctionExecutionResult( - content="4", - name=add_tool.name, - call_id=create_result.content[0].id, - is_error=False, + +@pytest.mark.asyncio +async def test_create_tools_with_thought(monkeypatch: pytest.MonkeyPatch) -> None: + def add(x: int, y: int) -> str: + return str(x + y) + + add_tool = FunctionTool(add, description="Add two numbers") + model = "llama3.2" + thought_content = "I'll use the add tool to calculate 2 + 2." + + async def _mock_chat(*args: Any, **kwargs: Any) -> ChatResponse: + return ChatResponse( + model=model, + done=True, + done_reason="tool_calls", + message=Message( + role="assistant", + content=thought_content, + tool_calls=[ + Message.ToolCall( + function=Message.ToolCall.Function( + name=add_tool.name, + arguments={"x": 2, "y": 2}, + ), + ), + ], + ), + prompt_eval_count=10, + eval_count=12, + ) + + monkeypatch.setattr(AsyncClient, "chat", _mock_chat) + client = OllamaChatCompletionClient(model=model) + + create_result = await client.create( + messages=[ + UserMessage(content="What is 2 + 2?", source="user"), + ], + tools=[add_tool], ) - stream = ollama_client.create_stream( + + assert isinstance(create_result.content, list) + assert len(create_result.content) > 0 + assert isinstance(create_result.content[0], FunctionCall) + assert create_result.content[0].name == add_tool.name + assert create_result.content[0].arguments == json.dumps({"x": 2, "y": 2}) + + assert create_result.thought == thought_content + + assert create_result.finish_reason == "function_calls" + assert create_result.usage is not None + assert create_result.usage.prompt_tokens == 10 + assert create_result.usage.completion_tokens == 12 + + +@pytest.mark.asyncio +async def test_create_stream_tools_with_thought(monkeypatch: pytest.MonkeyPatch) -> None: + def add(x: int, y: int) -> str: + return str(x + y) + + add_tool = FunctionTool(add, description="Add two numbers") + model = "llama3.2" + thought_content = "I'll use the add tool to calculate 2 + 2." + + async def _mock_chat(*args: Any, **kwargs: Any) -> AsyncGenerator[ChatResponse, None]: + assert "stream" in kwargs + assert kwargs["stream"] is True + + async def _mock_stream() -> AsyncGenerator[ChatResponse, None]: + thought_chunks = [thought_content[i : i + 10] for i in range(0, len(thought_content), 10)] + for chunk in thought_chunks: + yield ChatResponse( + model=model, + done=False, + message=Message( + role="assistant", + content=chunk, + ), + ) + + yield ChatResponse( + model=model, + done=True, + done_reason="tool_calls", + message=Message( + role="assistant", + tool_calls=[ + Message.ToolCall( + function=Message.ToolCall.Function( + name=add_tool.name, + arguments={"x": 2, "y": 2}, + ), + ), + ], + ), + prompt_eval_count=10, + eval_count=12, + ) + + return _mock_stream() + + monkeypatch.setattr(AsyncClient, "chat", _mock_chat) + client = OllamaChatCompletionClient(model=model) + + stream = client.create_stream( messages=[ - UserMessage( - content="What is 2 + 2? Use the add tool.", - source="user", - ), - AssistantMessage( - content=create_result.content, - source="assistant", - ), - FunctionExecutionResultMessage( - content=[execution_result], - ), + UserMessage(content="What is 2 + 2?", source="user"), ], + tools=[add_tool], ) - chunks = [] + + chunks: List[str | CreateResult] = [] async for chunk in stream: chunks.append(chunk) + assert len(chunks) > 0 - assert isinstance(chunks[-1], CreateResult) - create_result = chunks[-1] - assert isinstance(create_result.content, str) + + create_result = next((c for c in chunks if isinstance(c, CreateResult)), None) + assert create_result is not None + + assert isinstance(create_result.content, list) assert len(create_result.content) > 0 - assert create_result.finish_reason == "stop" + assert isinstance(create_result.content[0], FunctionCall) + assert create_result.content[0].name == add_tool.name + assert create_result.content[0].arguments == json.dumps({"x": 2, "y": 2}) + + assert create_result.thought == thought_content + + assert create_result.finish_reason == "function_calls" + assert create_result.usage is not None + assert create_result.usage.prompt_tokens == 10 + assert create_result.usage.completion_tokens == 12 + + +@pytest.mark.asyncio +async def test_llm_control_params(monkeypatch: pytest.MonkeyPatch) -> None: + model_name = "llama3.2" + + # Capture the kwargs passed to chat + chat_kwargs_captured: Dict[str, Any] = {} + + async def _mock_chat(*args: Any, **kwargs: Any) -> ChatResponse: + nonlocal chat_kwargs_captured + chat_kwargs_captured = kwargs + return ChatResponse( + model=model_name, + done=True, + done_reason="stop", + message=Message( + role="assistant", + content="Test response", + ), + ) + + monkeypatch.setattr(AsyncClient, "chat", _mock_chat) + + client_params: Dict[str, Any] = {"model": model_name, "temperature": 0.7, "top_p": 0.9, "frequency_penalty": 1.2} + + client = OllamaChatCompletionClient(**client_params) + + await client.create( + messages=[ + UserMessage(content="hi", source="user"), + ], + ) + + assert "options" in chat_kwargs_captured + assert isinstance(chat_kwargs_captured["options"], dict) + assert chat_kwargs_captured["options"]["temperature"] == 0.7 + assert chat_kwargs_captured["options"]["top_p"] == 0.9 + assert chat_kwargs_captured["options"]["frequency_penalty"] == 1.2 + + +@pytest.mark.asyncio +async def test_tool_choice_auto(monkeypatch: pytest.MonkeyPatch) -> None: + """Test tool_choice='auto' (default behavior)""" + + def add(x: int, y: int) -> str: + return str(x + y) + + def multiply(x: int, y: int) -> str: + return str(x * y) + + add_tool = FunctionTool(add, description="Add two numbers") + multiply_tool = FunctionTool(multiply, description="Multiply two numbers") + model = "llama3.2" + + # Capture the kwargs passed to chat + chat_kwargs_captured: Dict[str, Any] = {} + + async def _mock_chat(*args: Any, **kwargs: Any) -> ChatResponse: + nonlocal chat_kwargs_captured + chat_kwargs_captured = kwargs + return ChatResponse( + model=model, + done=True, + done_reason="stop", + message=Message( + role="assistant", + content="I'll use the add tool.", + tool_calls=[ + Message.ToolCall( + function=Message.ToolCall.Function( + name=add_tool.name, + arguments={"x": 2, "y": 3}, + ), + ), + ], + ), + prompt_eval_count=10, + eval_count=12, + ) + + monkeypatch.setattr(AsyncClient, "chat", _mock_chat) + + client = OllamaChatCompletionClient(model=model) + create_result = await client.create( + messages=[UserMessage(content="What is 2 + 3?", source="user")], + tools=[add_tool, multiply_tool], + tool_choice="auto", # Explicitly set to auto + ) + + # Verify that all tools are passed to the API when tool_choice is auto + assert "tools" in chat_kwargs_captured + assert chat_kwargs_captured["tools"] is not None + assert len(chat_kwargs_captured["tools"]) == 2 + + # Verify the response + assert isinstance(create_result.content, list) + assert len(create_result.content) > 0 + assert isinstance(create_result.content[0], FunctionCall) + assert create_result.content[0].name == add_tool.name + + +@pytest.mark.asyncio +async def test_tool_choice_none(monkeypatch: pytest.MonkeyPatch) -> None: + """Test tool_choice='none' - no tools should be passed to API""" + + def add(x: int, y: int) -> str: + return str(x + y) + + add_tool = FunctionTool(add, description="Add two numbers") + model = "llama3.2" + content_raw = "I cannot use tools, so I'll calculate manually: 2 + 3 = 5" + + # Capture the kwargs passed to chat + chat_kwargs_captured: Dict[str, Any] = {} + + async def _mock_chat(*args: Any, **kwargs: Any) -> ChatResponse: + nonlocal chat_kwargs_captured + chat_kwargs_captured = kwargs + return ChatResponse( + model=model, + done=True, + done_reason="stop", + message=Message( + role="assistant", + content=content_raw, + ), + prompt_eval_count=10, + eval_count=12, + ) + + monkeypatch.setattr(AsyncClient, "chat", _mock_chat) + + client = OllamaChatCompletionClient(model=model) + create_result = await client.create( + messages=[UserMessage(content="What is 2 + 3?", source="user")], + tools=[add_tool], + tool_choice="none", + ) + + # Verify that no tools are passed to the API when tool_choice is none + assert "tools" in chat_kwargs_captured + assert chat_kwargs_captured["tools"] is None + + # Verify the response is text content + assert isinstance(create_result.content, str) + assert create_result.content == content_raw + assert create_result.finish_reason == "stop" + + +@pytest.mark.asyncio +async def test_tool_choice_required(monkeypatch: pytest.MonkeyPatch) -> None: + """Test tool_choice='required' - tools must be provided and passed to API""" + + def add(x: int, y: int) -> str: + return str(x + y) + + def multiply(x: int, y: int) -> str: + return str(x * y) + + add_tool = FunctionTool(add, description="Add two numbers") + multiply_tool = FunctionTool(multiply, description="Multiply two numbers") + model = "llama3.2" + + # Capture the kwargs passed to chat + chat_kwargs_captured: Dict[str, Any] = {} + + async def _mock_chat(*args: Any, **kwargs: Any) -> ChatResponse: + nonlocal chat_kwargs_captured + chat_kwargs_captured = kwargs + return ChatResponse( + model=model, + done=True, + done_reason="tool_calls", + message=Message( + role="assistant", + tool_calls=[ + Message.ToolCall( + function=Message.ToolCall.Function( + name=add_tool.name, + arguments={"x": 2, "y": 3}, + ), + ), + ], + ), + prompt_eval_count=10, + eval_count=12, + ) + + monkeypatch.setattr(AsyncClient, "chat", _mock_chat) + + client = OllamaChatCompletionClient(model=model) + create_result = await client.create( + messages=[UserMessage(content="What is 2 + 3?", source="user")], + tools=[add_tool, multiply_tool], + tool_choice="required", + ) + + # Verify that all tools are passed to the API when tool_choice is required + assert "tools" in chat_kwargs_captured + assert chat_kwargs_captured["tools"] is not None + assert len(chat_kwargs_captured["tools"]) == 2 + + # Verify the response contains function calls + assert isinstance(create_result.content, list) + assert len(create_result.content) > 0 + assert isinstance(create_result.content[0], FunctionCall) + assert create_result.content[0].name == add_tool.name + assert create_result.finish_reason == "function_calls" + + +@pytest.mark.asyncio +async def test_tool_choice_required_no_tools_error() -> None: + """Test tool_choice='required' with no tools raises ValueError""" + model = "llama3.2" + client = OllamaChatCompletionClient(model=model) + + with pytest.raises(ValueError, match="tool_choice 'required' specified but no tools provided"): + await client.create( + messages=[UserMessage(content="What is 2 + 3?", source="user")], + tools=[], # No tools provided + tool_choice="required", + ) + + +@pytest.mark.asyncio +async def test_tool_choice_specific_tool(monkeypatch: pytest.MonkeyPatch) -> None: + """Test tool_choice with a specific tool - only that tool should be passed to API""" + + def add(x: int, y: int) -> str: + return str(x + y) + + def multiply(x: int, y: int) -> str: + return str(x * y) + + add_tool = FunctionTool(add, description="Add two numbers") + multiply_tool = FunctionTool(multiply, description="Multiply two numbers") + model = "llama3.2" + + # Capture the kwargs passed to chat + chat_kwargs_captured: Dict[str, Any] = {} + + async def _mock_chat(*args: Any, **kwargs: Any) -> ChatResponse: + nonlocal chat_kwargs_captured + chat_kwargs_captured = kwargs + return ChatResponse( + model=model, + done=True, + done_reason="tool_calls", + message=Message( + role="assistant", + tool_calls=[ + Message.ToolCall( + function=Message.ToolCall.Function( + name=add_tool.name, + arguments={"x": 2, "y": 3}, + ), + ), + ], + ), + prompt_eval_count=10, + eval_count=12, + ) + + monkeypatch.setattr(AsyncClient, "chat", _mock_chat) + + client = OllamaChatCompletionClient(model=model) + create_result = await client.create( + messages=[UserMessage(content="What is 2 + 3?", source="user")], + tools=[add_tool, multiply_tool], # Multiple tools available + tool_choice=add_tool, # But force specific tool + ) + + # Verify that only the specified tool is passed to the API + assert "tools" in chat_kwargs_captured + assert chat_kwargs_captured["tools"] is not None + assert len(chat_kwargs_captured["tools"]) == 1 + assert chat_kwargs_captured["tools"][0]["function"]["name"] == add_tool.name + + # Verify the response contains function calls + assert isinstance(create_result.content, list) + assert len(create_result.content) > 0 + assert isinstance(create_result.content[0], FunctionCall) + assert create_result.content[0].name == add_tool.name + assert create_result.finish_reason == "function_calls" + + +@pytest.mark.asyncio +async def test_tool_choice_stream_auto(monkeypatch: pytest.MonkeyPatch) -> None: + """Test tool_choice='auto' with streaming""" + + def add(x: int, y: int) -> str: + return str(x + y) + + add_tool = FunctionTool(add, description="Add two numbers") + model = "llama3.2" + content_raw = "I'll use the add tool." + + # Capture the kwargs passed to chat + chat_kwargs_captured: Dict[str, Any] = {} + + async def _mock_chat(*args: Any, **kwargs: Any) -> AsyncGenerator[ChatResponse, None]: + nonlocal chat_kwargs_captured + chat_kwargs_captured = kwargs + assert "stream" in kwargs + assert kwargs["stream"] is True + + async def _mock_stream() -> AsyncGenerator[ChatResponse, None]: + chunks = [content_raw[i : i + 5] for i in range(0, len(content_raw), 5)] + # Simulate streaming by yielding chunks of the response + for chunk in chunks[:-1]: + yield ChatResponse( + model=model, + done=False, + message=Message( + role="assistant", + content=chunk, + ), + ) + yield ChatResponse( + model=model, + done=True, + done_reason="tool_calls", + message=Message( + content=chunks[-1], + role="assistant", + tool_calls=[ + Message.ToolCall( + function=Message.ToolCall.Function( + name=add_tool.name, + arguments={"x": 2, "y": 3}, + ), + ), + ], + ), + prompt_eval_count=10, + eval_count=12, + ) + + return _mock_stream() + + monkeypatch.setattr(AsyncClient, "chat", _mock_chat) + + client = OllamaChatCompletionClient(model=model) + stream = client.create_stream( + messages=[UserMessage(content="What is 2 + 3?", source="user")], + tools=[add_tool], + tool_choice="auto", + ) + + chunks: List[str | CreateResult] = [] + async for chunk in stream: + chunks.append(chunk) + + # Verify that tools are passed to the API when tool_choice is auto + assert "tools" in chat_kwargs_captured + assert chat_kwargs_captured["tools"] is not None + assert len(chat_kwargs_captured["tools"]) == 1 + + # Verify the final result + assert len(chunks) > 0 + assert isinstance(chunks[-1], CreateResult) + assert isinstance(chunks[-1].content, list) + assert len(chunks[-1].content) > 0 + assert isinstance(chunks[-1].content[0], FunctionCall) + assert chunks[-1].content[0].name == add_tool.name + assert chunks[-1].finish_reason == "function_calls" + + +@pytest.mark.asyncio +async def test_tool_choice_stream_none(monkeypatch: pytest.MonkeyPatch) -> None: + """Test tool_choice='none' with streaming""" + + def add(x: int, y: int) -> str: + return str(x + y) + + add_tool = FunctionTool(add, description="Add two numbers") + model = "llama3.2" + content_raw = "I cannot use tools, so I'll calculate manually: 2 + 3 = 5" + + # Capture the kwargs passed to chat + chat_kwargs_captured: Dict[str, Any] = {} + + async def _mock_chat(*args: Any, **kwargs: Any) -> AsyncGenerator[ChatResponse, None]: + nonlocal chat_kwargs_captured + chat_kwargs_captured = kwargs + assert "stream" in kwargs + assert kwargs["stream"] is True + + async def _mock_stream() -> AsyncGenerator[ChatResponse, None]: + chunks = [content_raw[i : i + 10] for i in range(0, len(content_raw), 10)] + # Simulate streaming by yielding chunks of the response + for chunk in chunks[:-1]: + yield ChatResponse( + model=model, + done=False, + message=Message( + role="assistant", + content=chunk, + ), + ) + yield ChatResponse( + model=model, + done=True, + done_reason="stop", + message=Message( + role="assistant", + content=chunks[-1], + ), + prompt_eval_count=10, + eval_count=12, + ) + + return _mock_stream() + + monkeypatch.setattr(AsyncClient, "chat", _mock_chat) + + client = OllamaChatCompletionClient(model=model) + stream = client.create_stream( + messages=[UserMessage(content="What is 2 + 3?", source="user")], + tools=[add_tool], + tool_choice="none", + ) + + chunks: List[str | CreateResult] = [] + async for chunk in stream: + chunks.append(chunk) + + # Verify that no tools are passed to the API when tool_choice is none + assert "tools" in chat_kwargs_captured + assert chat_kwargs_captured["tools"] is None + + # Verify the final result is text content + assert len(chunks) > 0 + assert isinstance(chunks[-1], CreateResult) + assert isinstance(chunks[-1].content, str) + assert chunks[-1].content == content_raw + assert chunks[-1].finish_reason == "stop" + + +@pytest.mark.asyncio +async def test_tool_choice_default_behavior(monkeypatch: pytest.MonkeyPatch) -> None: + """Test that default behavior (no tool_choice specified) works like 'auto'""" + + def add(x: int, y: int) -> str: + return str(x + y) + + add_tool = FunctionTool(add, description="Add two numbers") + model = "llama3.2" + + # Capture the kwargs passed to chat + chat_kwargs_captured: Dict[str, Any] = {} + + async def _mock_chat(*args: Any, **kwargs: Any) -> ChatResponse: + nonlocal chat_kwargs_captured + chat_kwargs_captured = kwargs + return ChatResponse( + model=model, + done=True, + done_reason="stop", + message=Message( + role="assistant", + content="I'll use the add tool.", + tool_calls=[ + Message.ToolCall( + function=Message.ToolCall.Function( + name=add_tool.name, + arguments={"x": 2, "y": 3}, + ), + ), + ], + ), + prompt_eval_count=10, + eval_count=12, + ) + + monkeypatch.setattr(AsyncClient, "chat", _mock_chat) + + client = OllamaChatCompletionClient(model=model) + create_result = await client.create( + messages=[UserMessage(content="What is 2 + 3?", source="user")], + tools=[add_tool], + # tool_choice not specified - should default to "auto" + ) + + # Verify that tools are passed to the API by default (auto behavior) + assert "tools" in chat_kwargs_captured + assert chat_kwargs_captured["tools"] is not None + assert len(chat_kwargs_captured["tools"]) == 1 + + # Verify the response + assert isinstance(create_result.content, list) + assert len(create_result.content) > 0 + assert isinstance(create_result.content[0], FunctionCall) + assert create_result.content[0].name == add_tool.name diff --git a/python/packages/autogen-ext/tests/models/test_openai_model_client.py b/python/packages/autogen-ext/tests/models/test_openai_model_client.py index 7ae6aeee0796..58558cceb5f4 100644 --- a/python/packages/autogen-ext/tests/models/test_openai_model_client.py +++ b/python/packages/autogen-ext/tests/models/test_openai_model_client.py @@ -3,10 +3,12 @@ import logging import os from typing import Annotated, Any, AsyncGenerator, Dict, List, Literal, Tuple, TypeVar -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import httpx import pytest +from autogen_agentchat.agents import AssistantAgent +from autogen_agentchat.messages import MultiModalMessage from autogen_core import CancellationToken, FunctionCall, Image from autogen_core.models import ( AssistantMessage, @@ -24,18 +26,14 @@ from autogen_ext.models.openai import AzureOpenAIChatCompletionClient, OpenAIChatCompletionClient from autogen_ext.models.openai._model_info import resolve_model from autogen_ext.models.openai._openai_client import ( + BaseOpenAIChatCompletionClient, calculate_vision_tokens, convert_tools, to_oai_type, ) -from openai.resources.beta.chat.completions import ( # type: ignore - AsyncChatCompletionStreamManager as BetaAsyncChatCompletionStreamManager, # type: ignore -) - -# type: ignore -from openai.resources.beta.chat.completions import ( - AsyncCompletions as BetaAsyncCompletions, -) +from autogen_ext.models.openai._transformation import TransformerMap, get_transformer +from autogen_ext.models.openai._transformation.registry import _find_model_family # pyright: ignore[reportPrivateUsage] +from openai.lib.streaming.chat import AsyncChatCompletionStreamManager from openai.resources.chat.completions import AsyncCompletions from openai.types.chat.chat_completion import ChatCompletion, Choice from openai.types.chat.chat_completion_chunk import ( @@ -92,7 +90,7 @@ class MockChunkEvent(BaseModel): async def _mock_create_stream(*args: Any, **kwargs: Any) -> AsyncGenerator[ChatCompletionChunk, None]: - model = resolve_model(kwargs.get("model", "gpt-4o")) + model = resolve_model(kwargs.get("model", "gpt-4.1-nano")) mock_chunks_content = ["Hello", " Another Hello", " Yet Another Hello"] # The openai api implementations (OpenAI and Litellm) stream chunks of tokens @@ -164,7 +162,7 @@ async def _mock_create_stream(*args: Any, **kwargs: Any) -> AsyncGenerator[ChatC async def _mock_create(*args: Any, **kwargs: Any) -> ChatCompletion | AsyncGenerator[ChatCompletionChunk, None]: stream = kwargs.get("stream", False) - model = resolve_model(kwargs.get("model", "gpt-4o")) + model = resolve_model(kwargs.get("model", "gpt-4.1-nano")) if not stream: await asyncio.sleep(0.1) return ChatCompletion( @@ -183,7 +181,7 @@ async def _mock_create(*args: Any, **kwargs: Any) -> ChatCompletion | AsyncGener @pytest.mark.asyncio async def test_openai_chat_completion_client() -> None: - client = OpenAIChatCompletionClient(model="gpt-4o", api_key="api_key") + client = OpenAIChatCompletionClient(model="gpt-4.1-nano", api_key="api_key") assert client @@ -195,7 +193,7 @@ async def test_openai_chat_completion_client_with_gemini_model() -> None: @pytest.mark.asyncio async def test_openai_chat_completion_client_serialization() -> None: - client = OpenAIChatCompletionClient(model="gpt-4o", api_key="sk-password") + client = OpenAIChatCompletionClient(model="gpt-4.1-nano", api_key="sk-password") assert client config = client.dump_component() assert config @@ -271,6 +269,7 @@ async def test_openai_chat_completion_client_create_stream_with_usage( monkeypatch.setattr(AsyncCompletions, "create", _mock_create) client = OpenAIChatCompletionClient(model="gpt-4o", api_key="api_key") chunks: List[str | CreateResult] = [] + # Check that include_usage works when set via create_args with caplog.at_level(logging.INFO): async for chunk in client.create_stream( messages=[UserMessage(content="Hello", source="user")], @@ -291,6 +290,38 @@ async def test_openai_chat_completion_client_create_stream_with_usage( assert chunks[-1].content in caplog.text assert chunks[-1].usage == RequestUsage(prompt_tokens=3, completion_tokens=3) + chunks = [] + # Check that include_usage works when set via include_usage flag + with caplog.at_level(logging.INFO): + async for chunk in client.create_stream( + messages=[UserMessage(content="Hello", source="user")], + include_usage=True, + ): + chunks.append(chunk) + + assert "LLMStreamStart" in caplog.text + assert "LLMStreamEnd" in caplog.text + + assert chunks[0] == "Hello" + assert chunks[1] == " Another Hello" + assert chunks[2] == " Yet Another Hello" + assert isinstance(chunks[-1], CreateResult) + assert isinstance(chunks[-1].content, str) + assert chunks[-1].content == "Hello Another Hello Yet Another Hello" + assert chunks[-1].content in caplog.text + assert chunks[-1].usage == RequestUsage(prompt_tokens=3, completion_tokens=3) + + chunks = [] + # Check that setting both flags to different values raises an exception + + with pytest.raises(ValueError): + async for chunk in client.create_stream( + messages=[UserMessage(content="Hello", source="user")], + extra_create_args={"stream_options": {"include_usage": False}}, + include_usage=True, + ): + chunks.append(chunk) + @pytest.mark.asyncio async def test_openai_chat_completion_client_create_stream_no_usage_default(monkeypatch: pytest.MonkeyPatch) -> None: @@ -329,9 +360,39 @@ async def test_openai_chat_completion_client_create_stream_no_usage_explicit(mon assert chunks[0] == "Hello" assert chunks[1] == " Another Hello" assert chunks[2] == " Yet Another Hello" - assert isinstance(chunks[-1], CreateResult) - assert chunks[-1].content == "Hello Another Hello Yet Another Hello" - assert chunks[-1].usage == RequestUsage(prompt_tokens=0, completion_tokens=0) + + +@pytest.mark.asyncio +async def test_openai_chat_completion_client_none_usage(monkeypatch: pytest.MonkeyPatch) -> None: + """Test that completion_tokens and prompt_tokens handle None usage correctly. + + This test addresses issue #6352 where result.usage could be None, + causing TypeError in logging when trying to access completion_tokens. + """ + + async def _mock_create_with_none_usage(*args: Any, **kwargs: Any) -> ChatCompletion: + await asyncio.sleep(0.1) + # Create a ChatCompletion with None usage (which can happen in some API scenarios) + return ChatCompletion( + id="id", + choices=[ + Choice(finish_reason="stop", index=0, message=ChatCompletionMessage(content="Hello", role="assistant")) + ], + created=0, + model="gpt-4o", + object="chat.completion", + usage=None, # This is the scenario from the issue + ) + + monkeypatch.setattr(AsyncCompletions, "create", _mock_create_with_none_usage) + client = OpenAIChatCompletionClient(model="gpt-4o", api_key="api_key") + + # This should not raise a TypeError + result = await client.create(messages=[UserMessage(content="Hello", source="user")]) + + # Verify that the usage is correctly set to 0 when usage is None + assert result.usage.prompt_tokens == 0 + assert result.usage.completion_tokens == 0 @pytest.mark.asyncio @@ -464,7 +525,7 @@ async def run(self, args: MyArgs, cancellation_token: CancellationToken) -> MyRe @pytest.mark.asyncio async def test_json_mode(monkeypatch: pytest.MonkeyPatch) -> None: - model = "gpt-4o-2024-11-20" + model = "gpt-4.1-nano-2025-04-14" called_args = {} @@ -559,7 +620,7 @@ class AgentResponse(BaseModel): thoughts: str response: Literal["happy", "sad", "neutral"] - model = "gpt-4o-2024-11-20" + model = "gpt-4.1-nano-2025-04-14" called_args = {} @@ -651,7 +712,7 @@ class AgentResponse(BaseModel): thoughts: str response: Literal["happy", "sad", "neutral"] - model = "gpt-4o-2024-11-20" + model = "gpt-4.1-nano-2025-04-14" async def _mock_parse(*args: Any, **kwargs: Any) -> ParsedChatCompletion[AgentResponse]: return ParsedChatCompletion( @@ -677,7 +738,7 @@ async def _mock_parse(*args: Any, **kwargs: Any) -> ParsedChatCompletion[AgentRe usage=CompletionUsage(prompt_tokens=10, completion_tokens=5, total_tokens=0), ) - monkeypatch.setattr(BetaAsyncCompletions, "parse", _mock_parse) + monkeypatch.setattr(AsyncCompletions, "parse", _mock_parse) model_client = OpenAIChatCompletionClient( model=model, @@ -734,7 +795,7 @@ class AgentResponse(BaseModel): thoughts: str response: Literal["happy", "sad", "neutral"] - model = "gpt-4o-2024-11-20" + model = "gpt-4.1-nano-2025-04-14" async def _mock_parse(*args: Any, **kwargs: Any) -> ParsedChatCompletion[AgentResponse]: return ParsedChatCompletion( @@ -770,7 +831,7 @@ async def _mock_parse(*args: Any, **kwargs: Any) -> ParsedChatCompletion[AgentRe usage=CompletionUsage(prompt_tokens=10, completion_tokens=5, total_tokens=0), ) - monkeypatch.setattr(BetaAsyncCompletions, "parse", _mock_parse) + monkeypatch.setattr(AsyncCompletions, "parse", _mock_parse) model_client = OpenAIChatCompletionClient( model=model, @@ -810,7 +871,7 @@ class AgentResponse(BaseModel): chunked_content = [raw_content[i : i + 5] for i in range(0, len(raw_content), 5)] assert "".join(chunked_content) == raw_content - model = "gpt-4o-2024-11-20" + model = "gpt-4.1-nano-2025-04-14" mock_chunk_events = [ MockChunkEvent( type="chunk", @@ -844,7 +905,7 @@ async def _stream() -> AsyncGenerator[MockChunkEvent, None]: return _stream() # Mock the context manager __aenter__ method which returns the stream. - monkeypatch.setattr(BetaAsyncChatCompletionStreamManager, "__aenter__", _mock_create_stream) + monkeypatch.setattr(AsyncChatCompletionStreamManager, "__aenter__", _mock_create_stream) model_client = OpenAIChatCompletionClient( model=model, @@ -883,7 +944,7 @@ class AgentResponse(BaseModel): chunked_content = [raw_content[i : i + 5] for i in range(0, len(raw_content), 5)] assert "".join(chunked_content) == raw_content - model = "gpt-4o-2024-11-20" + model = "gpt-4.1-nano-2025-04-14" # generate the list of mock chunk content mock_chunk_events = [ @@ -954,7 +1015,7 @@ async def _stream() -> AsyncGenerator[MockChunkEvent, None]: return _stream() # Mock the context manager __aenter__ method which returns the stream. - monkeypatch.setattr(BetaAsyncChatCompletionStreamManager, "__aenter__", _mock_create_stream) + monkeypatch.setattr(AsyncChatCompletionStreamManager, "__aenter__", _mock_create_stream) model_client = OpenAIChatCompletionClient( model=model, @@ -1262,7 +1323,7 @@ async def _mock_create(*args: Any, **kwargs: Any) -> ChatCompletion | AsyncGener @pytest.mark.asyncio async def test_tool_calling(monkeypatch: pytest.MonkeyPatch) -> None: - model = "gpt-4o-2024-05-13" + model = "gpt-4.1-nano-2025-04-14" chat_completions = [ # Successful completion, single tool call ChatCompletion( @@ -1593,6 +1654,120 @@ async def _mock_create(*args: Any, **kwargs: Any) -> ChatCompletion | AsyncGener assert chunks[-1].thought == "Hello Another Hello Yet Another Hello" +@pytest.mark.asyncio +async def test_tool_calls_assistant_message_content_field(monkeypatch: pytest.MonkeyPatch) -> None: + """Test that AssistantMessage with tool calls includes required content field. + + This test addresses the issue where AssistantMessage with tool calls but no thought + was missing the required 'content' field, causing OpenAI API UnprocessableEntityError(422). + """ + # Create a tool call for testing + tool_calls = [ + FunctionCall(id="call_1", name="increment_number", arguments='{"number": 5}'), + FunctionCall(id="call_2", name="increment_number", arguments='{"number": 6}'), + ] + + # Mock response for tool calls + chat_completion = ChatCompletion( + id="id1", + choices=[ + Choice( + finish_reason="stop", + index=0, + message=ChatCompletionMessage( + role="assistant", + content="Done", + ), + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion", + usage=CompletionUsage(completion_tokens=10, prompt_tokens=5, total_tokens=15), + ) + + client = OpenAIChatCompletionClient(model="gpt-4o", api_key="test") + mock_create = AsyncMock(return_value=chat_completion) + + # Test AssistantMessage with tool calls but no thought + assistant_message_no_thought = AssistantMessage( + content=tool_calls, + source="assistant", + thought=None, # No thought - this was causing the issue + ) + + with monkeypatch.context() as mp: + mp.setattr(client._client.chat.completions, "create", mock_create) # type: ignore[reportPrivateUsage] + + await client.create( + messages=[ + UserMessage(content="Please increment these numbers", source="user"), + assistant_message_no_thought, + ] + ) + + # Verify the API was called and check the messages sent + mock_create.assert_called_once() + call_args = mock_create.call_args + + # Extract the messages from the API call + messages = call_args.kwargs["messages"] + + # Find the assistant message in the API call + assistant_messages = [msg for msg in messages if msg["role"] == "assistant"] + assert len(assistant_messages) == 1 + + assistant_msg = assistant_messages[0] + + # Verify all required fields are present + assert "role" in assistant_msg + assert "tool_calls" in assistant_msg + assert "content" in assistant_msg # This was missing before the fix + + # Verify field values + assert assistant_msg["role"] == "assistant" + assert assistant_msg["content"] is None # Should be null for tools without thought + assert len(assistant_msg["tool_calls"]) == 2 + + # Test AssistantMessage with tool calls AND thought + assistant_message_with_thought = AssistantMessage( + content=tool_calls, source="assistant", thought="I need to increment these numbers." + ) + + mock_create.reset_mock() # Reset for second test + + with monkeypatch.context() as mp: + mp.setattr(client._client.chat.completions, "create", mock_create) # type: ignore[reportPrivateUsage] + + await client.create( + messages=[ + UserMessage(content="Please increment these numbers", source="user"), + assistant_message_with_thought, + ] + ) + + # Verify the API was called for the second test + mock_create.assert_called_once() + call_args = mock_create.call_args + + # Extract the messages from the API call + messages = call_args.kwargs["messages"] + + # Find the assistant message in the API call + assistant_messages = [msg for msg in messages if msg["role"] == "assistant"] + assert len(assistant_messages) == 1 + + assistant_msg_with_thought = assistant_messages[0] + + # Should have both tool_calls and content with thought text + assert "role" in assistant_msg_with_thought + assert "tool_calls" in assistant_msg_with_thought + assert "content" in assistant_msg_with_thought + assert assistant_msg_with_thought["role"] == "assistant" + assert assistant_msg_with_thought["content"] == "I need to increment these numbers." + assert len(assistant_msg_with_thought["tool_calls"]) == 2 + + @pytest.fixture() def openai_client(request: pytest.FixtureRequest) -> OpenAIChatCompletionClient: model = request.node.callspec.params["model"] # type: ignore @@ -1601,6 +1776,10 @@ def openai_client(request: pytest.FixtureRequest) -> OpenAIChatCompletionClient: api_key = os.getenv("GEMINI_API_KEY") if not api_key: pytest.skip("GEMINI_API_KEY not found in environment variables") + elif model.startswith("claude"): + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not found in environment variables") else: api_key = os.getenv("OPENAI_API_KEY") if not api_key: @@ -1615,7 +1794,7 @@ def openai_client(request: pytest.FixtureRequest) -> OpenAIChatCompletionClient: @pytest.mark.asyncio @pytest.mark.parametrize( "model", - ["gpt-4o-mini", "gemini-1.5-flash"], + ["gpt-4.1-nano", "gemini-1.5-flash", "claude-3-5-haiku-20241022"], ) async def test_model_client_basic_completion(model: str, openai_client: OpenAIChatCompletionClient) -> None: # Test basic completion @@ -1632,13 +1811,15 @@ async def test_model_client_basic_completion(model: str, openai_client: OpenAICh @pytest.mark.asyncio @pytest.mark.parametrize( "model", - ["gpt-4o-mini", "gemini-1.5-flash"], + ["gpt-4.1-nano", "gemini-1.5-flash", "claude-3-5-haiku-20241022"], ) async def test_model_client_with_function_calling(model: str, openai_client: OpenAIChatCompletionClient) -> None: # Test tool calling pass_tool = FunctionTool(_pass_function, name="pass_tool", description="pass session.") fail_tool = FunctionTool(_fail_function, name="fail_tool", description="fail session.") - messages: List[LLMMessage] = [UserMessage(content="Call the pass tool with input 'task'", source="user")] + messages: List[LLMMessage] = [ + UserMessage(content="Call the pass tool with input 'task' summarize the result.", source="user") + ] create_result = await openai_client.create(messages=messages, tools=[pass_tool, fail_tool]) assert isinstance(create_result.content, list) assert len(create_result.content) == 1 @@ -1669,7 +1850,8 @@ async def test_model_client_with_function_calling(model: str, openai_client: Ope # Test parallel tool calling messages = [ UserMessage( - content="Call both the pass tool with input 'task' and the fail tool also with input 'task'", source="user" + content="Call both the pass tool with input 'task' and the fail tool also with input 'task' and summarize the result", + source="user", ) ] create_result = await openai_client.create(messages=messages, tools=[pass_tool, fail_tool]) @@ -1706,7 +1888,7 @@ async def test_model_client_with_function_calling(model: str, openai_client: Ope @pytest.mark.asyncio @pytest.mark.parametrize( "model", - ["gpt-4o-mini", "gemini-1.5-flash"], + ["gpt-4.1-nano", "gemini-1.5-flash"], ) async def test_openai_structured_output_using_response_format( model: str, openai_client: OpenAIChatCompletionClient @@ -1739,7 +1921,7 @@ class AgentResponse(BaseModel): @pytest.mark.asyncio @pytest.mark.parametrize( "model", - ["gpt-4o-mini", "gemini-1.5-flash"], + ["gpt-4.1-nano", "gemini-1.5-flash"], ) async def test_openai_structured_output(model: str, openai_client: OpenAIChatCompletionClient) -> None: class AgentResponse(BaseModel): @@ -1759,7 +1941,7 @@ class AgentResponse(BaseModel): @pytest.mark.asyncio @pytest.mark.parametrize( "model", - ["gpt-4o-mini", "gemini-1.5-flash"], + ["gpt-4.1-nano", "gemini-1.5-flash"], ) async def test_openai_structured_output_with_streaming(model: str, openai_client: OpenAIChatCompletionClient) -> None: class AgentResponse(BaseModel): @@ -1785,7 +1967,7 @@ class AgentResponse(BaseModel): @pytest.mark.parametrize( "model", [ - "gpt-4o-mini", + "gpt-4.1-nano", # "gemini-1.5-flash", # Gemini models do not support structured output with tool calls from model client. ], ) @@ -1843,7 +2025,7 @@ def sentiment_analysis(text: str) -> str: @pytest.mark.parametrize( "model", [ - "gpt-4o-mini", + "gpt-4.1-nano", # "gemini-1.5-flash", # Gemini models do not support structured output with tool calls from model client. ], ) @@ -2058,4 +2240,1016 @@ async def test_add_name_prefixes(monkeypatch: pytest.MonkeyPatch) -> None: assert str(converted_mm["content"][0]["text"]) == "Adam said:\n" + str(oai_mm["content"][0]["text"]) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", + [ + "gpt-4.1-nano", + "gemini-1.5-flash", + "claude-3-5-haiku-20241022", + ], +) +async def test_muliple_system_message(model: str, openai_client: OpenAIChatCompletionClient) -> None: + """Test multiple system messages in a single request.""" + + # Test multiple system messages + messages: List[LLMMessage] = [ + SystemMessage(content="When you say anything Start with 'FOO'"), + SystemMessage(content="When you say anything End with 'BAR'"), + UserMessage(content="Just say '.'", source="user"), + ] + + result = await openai_client.create(messages=messages) + result_content = result.content + assert isinstance(result_content, str) + result_content = result_content.strip() + assert result_content[:3] == "FOO" + assert result_content[-3:] == "BAR" + + +@pytest.mark.asyncio +async def test_system_message_merge_with_continuous_system_messages_models() -> None: + """Tests that system messages are merged correctly for Gemini models.""" + # Create a mock client + mock_client = MagicMock() + client = BaseOpenAIChatCompletionClient( + client=mock_client, + create_args={"model": "gemini-1.5-flash"}, + model_info={ + "vision": False, + "function_calling": False, + "json_output": False, + "family": "unknown", + "structured_output": False, + "multiple_system_messages": False, + }, + ) + + # Create two system messages + messages: List[LLMMessage] = [ + SystemMessage(content="I am system message 1"), + SystemMessage(content="I am system message 2"), + UserMessage(content="Hello", source="user"), + ] + + # Process the messages + # pylint: disable=protected-access + # The method is protected, but we need to test it + create_params = client._process_create_args( # pyright: ignore[reportPrivateUsage] + messages=messages, + tools=[], + json_output=None, + extra_create_args={}, + tool_choice="none", + ) + + # Extract the actual messages from the result + oai_messages = create_params.messages + + # Check that there is only one system message and it contains the merged content + system_messages = [msg for msg in oai_messages if msg["role"] == "system"] + assert len(system_messages) == 1 + assert system_messages[0]["content"] == "I am system message 1\nI am system message 2" + + # Check that the user message is preserved + user_messages = [msg for msg in oai_messages if msg["role"] == "user"] + assert len(user_messages) == 1 + assert user_messages[0]["content"] == "Hello" + + +@pytest.mark.asyncio +async def test_system_message_merge_with_non_continuous_messages() -> None: + """Tests that an error is raised when non-continuous system messages are provided.""" + # Create a mock client + mock_client = MagicMock() + client = BaseOpenAIChatCompletionClient( + client=mock_client, + create_args={"model": "gemini-1.5-flash"}, + model_info={ + "vision": False, + "function_calling": False, + "json_output": False, + "family": "unknown", + "structured_output": False, + "multiple_system_messages": False, + }, + ) + + # Create non-continuous system messages + messages: List[LLMMessage] = [ + SystemMessage(content="I am system message 1"), + UserMessage(content="Hello", source="user"), + SystemMessage(content="I am system message 2"), + ] + + # Process should raise ValueError + with pytest.raises(ValueError, match="Multiple and Not continuous system messages are not supported"): + # pylint: disable=protected-access + # The method is protected, but we need to test it + client._process_create_args( # pyright: ignore[reportPrivateUsage] + messages=messages, + tools=[], + json_output=None, + extra_create_args={}, + tool_choice="none", + ) + + +@pytest.mark.asyncio +async def test_system_message_not_merged_for_multiple_system_messages_true() -> None: + """Tests that system messages aren't modified for non-Gemini models.""" + # Create a mock client + mock_client = MagicMock() + client = BaseOpenAIChatCompletionClient( + client=mock_client, + create_args={"model": "gpt-4.1-nano"}, + model_info={ + "vision": False, + "function_calling": False, + "json_output": False, + "family": "unknown", + "structured_output": False, + "multiple_system_messages": True, + }, + ) + + # Create two system messages + messages: List[LLMMessage] = [ + SystemMessage(content="I am system message 1"), + SystemMessage(content="I am system message 2"), + UserMessage(content="Hello", source="user"), + ] + + # Process the messages + # pylint: disable=protected-access + # The method is protected, but we need to test it + create_params = client._process_create_args( # pyright: ignore[reportPrivateUsage] + messages=messages, + tools=[], + json_output=None, + extra_create_args={}, + tool_choice="none", + ) + + # Extract the actual messages from the result + oai_messages = create_params.messages + + # Check that there are two system messages preserved + system_messages = [msg for msg in oai_messages if msg["role"] == "system"] + assert len(system_messages) == 2 + assert system_messages[0]["content"] == "I am system message 1" + assert system_messages[1]["content"] == "I am system message 2" + + +@pytest.mark.asyncio +async def test_no_system_messages_for_gemini_model() -> None: + """Tests behavior when no system messages are provided to a Gemini model.""" + # Create a mock client + mock_client = MagicMock() + client = BaseOpenAIChatCompletionClient( + client=mock_client, + create_args={"model": "gemini-1.5-flash"}, + model_info={ + "vision": False, + "function_calling": False, + "json_output": False, + "family": "unknown", + "structured_output": False, + }, + ) + + # Create messages with no system message + messages: List[LLMMessage] = [ + UserMessage(content="Hello", source="user"), + AssistantMessage(content="Hi there", source="assistant"), + ] + + # Process the messages + # pylint: disable=protected-access + # The method is protected, but we need to test it + create_params = client._process_create_args( # pyright: ignore[reportPrivateUsage] + messages=messages, + tools=[], + json_output=None, + extra_create_args={}, + tool_choice="none", + ) + + # Extract the actual messages from the result + oai_messages = create_params.messages + + # Check that there are no system messages + system_messages = [msg for msg in oai_messages if msg["role"] == "system"] + assert len(system_messages) == 0 + + # Check that other messages are preserved + user_messages = [msg for msg in oai_messages if msg["role"] == "user"] + assistant_messages = [msg for msg in oai_messages if msg["role"] == "assistant"] + assert len(user_messages) == 1 + assert len(assistant_messages) == 1 + + +@pytest.mark.asyncio +async def test_single_system_message_for_gemini_model() -> None: + """Tests that a single system message is preserved for Gemini models.""" + # Create a mock client + mock_client = MagicMock() + client = BaseOpenAIChatCompletionClient( + client=mock_client, + create_args={"model": "gemini-1.5-flash"}, + model_info={ + "vision": False, + "function_calling": False, + "json_output": False, + "family": "unknown", + "structured_output": False, + }, + ) + + # Create messages with a single system message + messages: List[LLMMessage] = [ + SystemMessage(content="I am the only system message"), + UserMessage(content="Hello", source="user"), + ] + + # Process the messages + # pylint: disable=protected-access + # The method is protected, but we need to test it + create_params = client._process_create_args( # pyright: ignore[reportPrivateUsage] + messages=messages, + tools=[], + json_output=None, + extra_create_args={}, + tool_choice="auto", + ) + + # Extract the actual messages from the result + oai_messages = create_params.messages + + # Check that there is exactly one system message with the correct content + system_messages = [msg for msg in oai_messages if msg["role"] == "system"] + assert len(system_messages) == 1 + assert system_messages[0]["content"] == "I am the only system message" + + +def noop(input: str) -> str: + return "done" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["gemini-1.5-flash"]) +async def test_empty_assistant_content_with_gemini(model: str, openai_client: OpenAIChatCompletionClient) -> None: + # Test tool calling + tool = FunctionTool(noop, name="noop", description="No-op tool") + messages: List[LLMMessage] = [UserMessage(content="Call noop", source="user")] + result = await openai_client.create(messages=messages, tools=[tool]) + assert isinstance(result.content, list) + tool_call = result.content[0] + assert isinstance(tool_call, FunctionCall) + + # reply with empty string as thought (== content) + messages.append(AssistantMessage(content=result.content, thought="", source="assistant")) + messages.append( + FunctionExecutionResultMessage( + content=[ + FunctionExecutionResult( + content="done", + call_id=tool_call.id, + is_error=False, + name=tool_call.name, + ) + ] + ) + ) + + # This will crash if _set_empty_to_whitespace is not applied to "thought" + result = await openai_client.create(messages=messages) + assert isinstance(result.content, str) + assert result.content.strip() != "" or result.content == " " + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", + [ + "gpt-4.1-nano", + "gemini-1.5-flash", + "claude-3-5-haiku-20241022", + ], +) +async def test_empty_assistant_content_string_with_some_model( + model: str, openai_client: OpenAIChatCompletionClient +) -> None: + # message: assistant is response empty content + messages: list[LLMMessage] = [ + UserMessage(content="Say something", source="user"), + AssistantMessage(content="test", source="assistant"), + UserMessage(content="", source="user"), + ] + + # This will crash if _set_empty_to_whitespace is not applied to "content" + result = await openai_client.create(messages=messages) + assert isinstance(result.content, str) + + +def test_openai_model_registry_find_well() -> None: + model = "gpt-4o" + client1 = OpenAIChatCompletionClient(model=model, api_key="test") + client2 = OpenAIChatCompletionClient( + model=model, + model_info={ + "vision": False, + "function_calling": False, + "json_output": False, + "structured_output": False, + "family": ModelFamily.UNKNOWN, + }, + api_key="test", + ) + + def get_regitered_transformer(client: OpenAIChatCompletionClient) -> TransformerMap: + model_name = client._create_args["model"] # pyright: ignore[reportPrivateUsage] + model_family = client.model_info["family"] + return get_transformer("openai", model_name, model_family) + + assert get_regitered_transformer(client1) == get_regitered_transformer(client2) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", + [ + "gpt-4.1-nano", + ], +) +async def test_openai_model_unknown_message_type(model: str, openai_client: OpenAIChatCompletionClient) -> None: + class WrongMessage: + content = "foo" + source = "bar" + + messages: List[WrongMessage] = [WrongMessage()] + with pytest.raises(ValueError, match="Unknown message type"): + await openai_client.create(messages=messages) # type: ignore[arg-type] # pyright: ignore[reportArgumentType] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", + [ + "claude-3-5-haiku-20241022", + ], +) +async def test_claude_trailing_whitespace_at_last_assistant_content( + model: str, openai_client: OpenAIChatCompletionClient +) -> None: + messages: list[LLMMessage] = [ + UserMessage(content="foo", source="user"), + UserMessage(content="bar", source="user"), + AssistantMessage(content="foobar ", source="assistant"), + ] + + result = await openai_client.create(messages=messages) + assert isinstance(result.content, str) + + +def test_rstrip_railing_whitespace_at_last_assistant_content() -> None: + messages: list[LLMMessage] = [ + UserMessage(content="foo", source="user"), + UserMessage(content="bar", source="user"), + AssistantMessage(content="foobar ", source="assistant"), + ] + + # This will crash if _rstrip_railing_whitespace_at_last_assistant_content is not applied to "content" + dummy_client = OpenAIChatCompletionClient(model="claude-3-5-haiku-20241022", api_key="dummy-key") + result = dummy_client._rstrip_last_assistant_message(messages) # pyright: ignore[reportPrivateUsage] + + assert isinstance(result[-1].content, str) + assert result[-1].content == "foobar" + + +def test_find_model_family() -> None: + assert _find_model_family("openai", "gpt-4") == ModelFamily.GPT_4 + assert _find_model_family("openai", "gpt-4-latest") == ModelFamily.GPT_4 + assert _find_model_family("openai", "gpt-4o") == ModelFamily.GPT_4O + assert _find_model_family("openai", "gemini-2.0-flash") == ModelFamily.GEMINI_2_0_FLASH + assert _find_model_family("openai", "claude-3-5-haiku-20241022") == ModelFamily.CLAUDE_3_5_HAIKU + assert _find_model_family("openai", "error") == ModelFamily.UNKNOWN + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", + [ + "gpt-4.1-nano", + "gemini-1.5-flash", + "claude-3-5-haiku-20241022", + ], +) +async def test_multimodal_message_test( + model: str, openai_client: OpenAIChatCompletionClient, monkeypatch: pytest.MonkeyPatch +) -> None: + # Test that the multimodal message is converted to the correct format + img = Image.from_base64( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC" + ) + multi_modal_message = MultiModalMessage(content=["Can you describe the content of this image?", img], source="user") + + ocr_agent = AssistantAgent( + name="ocr_agent", model_client=openai_client, system_message="""You are a helpful agent.""" + ) + _ = await ocr_agent.run(task=multi_modal_message) + + +@pytest.mark.asyncio +async def test_mistral_remove_name() -> None: + # Test that the name pramaeter is removed from the message + # when the model is Mistral + message = UserMessage(content="foo", source="user") + params = to_oai_type(message, prepend_name=False, model="mistral-7b", model_family=ModelFamily.MISTRAL) + assert ("name" in params[0]) is False + + # when the model is gpt-4o, the name parameter is not removed + params = to_oai_type(message, prepend_name=False, model="gpt-4o", model_family=ModelFamily.GPT_4O) + assert ("name" in params[0]) is True + + +@pytest.mark.asyncio +async def test_include_name_in_message() -> None: + """Test that include_name_in_message parameter controls the name field.""" + + # Test with UserMessage + user_message = UserMessage(content="Hello, I am from Seattle.", source="Adam") + + # Test with include_name_in_message=True (default) + result_with_name = to_oai_type(user_message, include_name_in_message=True)[0] + assert "name" in result_with_name + assert result_with_name["name"] == "Adam" # type: ignore[typeddict-item] + assert result_with_name["role"] == "user" + assert result_with_name["content"] == "Hello, I am from Seattle." + + # Test with include_name_in_message=False + result_without_name = to_oai_type(user_message, include_name_in_message=False)[0] + assert "name" not in result_without_name + assert result_without_name["role"] == "user" + assert result_without_name["content"] == "Hello, I am from Seattle." + + # Test with AssistantMessage (should not have name field regardless) + assistant_message = AssistantMessage(content="Hello, how can I help you?", source="Assistant") + + # Test with include_name_in_message=True + result_assistant_with_name = to_oai_type(assistant_message, include_name_in_message=True)[0] + assert "name" not in result_assistant_with_name + assert result_assistant_with_name["role"] == "assistant" + + # Test with include_name_in_message=False + result_assistant_without_name = to_oai_type(assistant_message, include_name_in_message=False)[0] + assert "name" not in result_assistant_without_name + assert result_assistant_without_name["role"] == "assistant" + + # Test with SystemMessage (should not have name field regardless) + system_message = SystemMessage(content="You are a helpful assistant.") + result_system_with_name = to_oai_type(system_message, include_name_in_message=True)[0] + result_system_without_name = to_oai_type(system_message, include_name_in_message=False)[0] + assert "name" not in result_system_with_name + assert "name" not in result_system_without_name + assert result_system_with_name["role"] == "system" + assert result_system_without_name["role"] == "system" + + # Test default behavior (should include name when parameter not specified) + result_default = to_oai_type(user_message)[0] # include_name_in_message defaults to True + assert "name" in result_default + assert result_default["name"] == "Adam" # type: ignore[typeddict-item] + + +@pytest.mark.asyncio +async def test_include_name_with_different_models() -> None: + """Test that include_name_in_message works with different model families.""" + + user_message = UserMessage(content="Hello", source="User") + + # Test with GPT-4o model (normally includes name) + result_gpt4o_with_name = to_oai_type( + user_message, model="gpt-4o", model_family=ModelFamily.GPT_4O, include_name_in_message=True + )[0] + result_gpt4o_without_name = to_oai_type( + user_message, model="gpt-4o", model_family=ModelFamily.GPT_4O, include_name_in_message=False + )[0] + + assert "name" in result_gpt4o_with_name + assert "name" not in result_gpt4o_without_name + + # Test with Mistral model (normally excludes name, but should still respect the parameter) + result_mistral_with_name = to_oai_type( + user_message, model="mistral-7b", model_family=ModelFamily.MISTRAL, include_name_in_message=True + )[0] + result_mistral_without_name = to_oai_type( + user_message, model="mistral-7b", model_family=ModelFamily.MISTRAL, include_name_in_message=False + )[0] + + # Note: Mistral transformers are specifically built without _set_name, so they won't have name regardless + # But our parameter still controls the behavior consistently + assert "name" not in result_mistral_with_name # Mistral design excludes names + assert "name" not in result_mistral_without_name + + # Test with unknown model (uses default transformer) + result_unknown_with_name = to_oai_type( + user_message, model="some-custom-model", model_family=ModelFamily.UNKNOWN, include_name_in_message=True + )[0] + result_unknown_without_name = to_oai_type( + user_message, model="some-custom-model", model_family=ModelFamily.UNKNOWN, include_name_in_message=False + )[0] + + assert "name" in result_unknown_with_name + assert "name" not in result_unknown_without_name + + +@pytest.mark.asyncio +async def test_mock_tool_choice_specific_tool(monkeypatch: pytest.MonkeyPatch) -> None: + """Test tool_choice parameter with a specific tool using mocks.""" + + def _pass_function(input: str) -> str: + """Simple passthrough function.""" + return f"Processed: {input}" + + def _add_numbers(a: int, b: int) -> int: + """Add two numbers together.""" + return a + b + + model = "gpt-4o" + + # Mock successful completion with specific tool call + chat_completion = ChatCompletion( + id="id1", + choices=[ + Choice( + finish_reason="tool_calls", + index=0, + message=ChatCompletionMessage( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="1", + type="function", + function=Function( + name="_pass_function", + arguments=json.dumps({"input": "hello"}), + ), + ) + ], + ), + ) + ], + created=1234567890, + model=model, + object="chat.completion", + usage=CompletionUsage(completion_tokens=10, prompt_tokens=5, total_tokens=15), + ) + + client = OpenAIChatCompletionClient(model=model, api_key="test") + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="_pass_function") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="_add_numbers") + + # Create mock for the chat completions create method + mock_create = AsyncMock(return_value=chat_completion) + + with monkeypatch.context() as mp: + mp.setattr(client._client.chat.completions, "create", mock_create) # type: ignore[reportPrivateUsage] + + _ = await client.create( + messages=[UserMessage(content="Process 'hello'", source="user")], + tools=[pass_tool, add_tool], + tool_choice=pass_tool, # Force use of specific tool + ) + + # Verify the correct API call was made + mock_create.assert_called_once() + call_args = mock_create.call_args + + # Check that tool_choice was set correctly + assert "tool_choice" in call_args.kwargs + assert call_args.kwargs["tool_choice"] == {"type": "function", "function": {"name": "_pass_function"}} + + +@pytest.mark.asyncio +async def test_mock_tool_choice_auto(monkeypatch: pytest.MonkeyPatch) -> None: + """Test tool_choice parameter with 'auto' setting using mocks.""" + + def _pass_function(input: str) -> str: + """Simple passthrough function.""" + return f"Processed: {input}" + + def _add_numbers(a: int, b: int) -> int: + """Add two numbers together.""" + return a + b + + model = "gpt-4o" + + # Mock successful completion + chat_completion = ChatCompletion( + id="id1", + choices=[ + Choice( + finish_reason="tool_calls", + index=0, + message=ChatCompletionMessage( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="1", + type="function", + function=Function( + name="_add_numbers", + arguments=json.dumps({"a": 1, "b": 2}), + ), + ) + ], + ), + ) + ], + created=1234567890, + model=model, + object="chat.completion", + usage=CompletionUsage(completion_tokens=10, prompt_tokens=5, total_tokens=15), + ) + + client = OpenAIChatCompletionClient(model=model, api_key="test") + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="_pass_function") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="_add_numbers") + + # Create mock for the chat completions create method + mock_create = AsyncMock(return_value=chat_completion) + + with monkeypatch.context() as mp: + mp.setattr(client._client.chat.completions, "create", mock_create) # type: ignore[reportPrivateUsage] + + await client.create( + messages=[UserMessage(content="Add 1 and 2", source="user")], + tools=[pass_tool, add_tool], + tool_choice="auto", # Let model choose + ) + + # Verify the correct API call was made + mock_create.assert_called_once() + call_args = mock_create.call_args + + # Check that tool_choice was set correctly + assert "tool_choice" in call_args.kwargs + assert call_args.kwargs["tool_choice"] == "auto" + + +@pytest.mark.asyncio +async def test_mock_tool_choice_none(monkeypatch: pytest.MonkeyPatch) -> None: + """Test tool_choice parameter with None setting using mocks.""" + + def _pass_function(input: str) -> str: + """Simple passthrough function.""" + return f"Processed: {input}" + + model = "gpt-4o" + + # Mock successful completion + chat_completion = ChatCompletion( + id="id1", + choices=[ + Choice( + finish_reason="stop", + index=0, + message=ChatCompletionMessage( + role="assistant", + content="I can help you with that!", + tool_calls=None, + ), + ) + ], + created=1234567890, + model=model, + object="chat.completion", + usage=CompletionUsage(completion_tokens=10, prompt_tokens=5, total_tokens=15), + ) + + client = OpenAIChatCompletionClient(model=model, api_key="test") + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="_pass_function") + + # Create mock for the chat completions create method + mock_create = AsyncMock(return_value=chat_completion) + + with monkeypatch.context() as mp: + mp.setattr(client._client.chat.completions, "create", mock_create) # type: ignore[reportPrivateUsage] + + await client.create( + messages=[UserMessage(content="Hello there", source="user")], + tools=[pass_tool], + tool_choice="none", + ) + + # Verify the correct API call was made + mock_create.assert_called_once() + call_args = mock_create.call_args + + # Check that tool_choice was set to "none" (disabling tool usage) + assert "tool_choice" in call_args.kwargs + assert call_args.kwargs["tool_choice"] == "none" + + +@pytest.mark.asyncio +async def test_mock_tool_choice_validation_error() -> None: + """Test tool_choice validation with invalid tool reference.""" + + def _pass_function(input: str) -> str: + """Simple passthrough function.""" + return f"Processed: {input}" + + def _add_numbers(a: int, b: int) -> int: + """Add two numbers together.""" + return a + b + + def _different_function(text: str) -> str: + """Different function.""" + return text + + client = OpenAIChatCompletionClient(model="gpt-4o", api_key="test") + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="_pass_function") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="_add_numbers") + different_tool = FunctionTool(_different_function, description="Different tool", name="_different_function") + + messages = [UserMessage(content="Hello there", source="user")] + + # Test with a tool that's not in the tools list + with pytest.raises( + ValueError, match="tool_choice references '_different_function' but it's not in the provided tools" + ): + await client.create( + messages=messages, + tools=[pass_tool, add_tool], + tool_choice=different_tool, # This tool is not in the tools list + ) + + +@pytest.mark.asyncio +async def test_mock_tool_choice_required(monkeypatch: pytest.MonkeyPatch) -> None: + """Test tool_choice parameter with 'required' setting using mocks.""" + + def _pass_function(input: str) -> str: + """Simple passthrough function.""" + return f"Processed: {input}" + + def _add_numbers(a: int, b: int) -> int: + """Add two numbers together.""" + return a + b + + model = "gpt-4o" + + # Mock successful completion with tool calls (required forces tool usage) + chat_completion = ChatCompletion( + id="id1", + choices=[ + Choice( + finish_reason="tool_calls", + index=0, + message=ChatCompletionMessage( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="1", + type="function", + function=Function( + name="_pass_function", + arguments=json.dumps({"input": "hello"}), + ), + ) + ], + ), + ) + ], + created=1234567890, + model=model, + object="chat.completion", + usage=CompletionUsage(completion_tokens=10, prompt_tokens=5, total_tokens=15), + ) + + client = OpenAIChatCompletionClient(model=model, api_key="test") + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="_pass_function") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="_add_numbers") + + # Create mock for the chat completions create method + mock_create = AsyncMock(return_value=chat_completion) + + with monkeypatch.context() as mp: + mp.setattr(client._client.chat.completions, "create", mock_create) # type: ignore[reportPrivateUsage] + + await client.create( + messages=[UserMessage(content="Process some text", source="user")], + tools=[pass_tool, add_tool], + tool_choice="required", # Force tool usage + ) + + # Verify the correct API call was made + mock_create.assert_called_once() + call_args = mock_create.call_args + + # Check that tool_choice was set correctly + assert "tool_choice" in call_args.kwargs + assert call_args.kwargs["tool_choice"] == "required" + + +# Integration tests for tool_choice using the actual OpenAI API +@pytest.mark.asyncio +async def test_openai_tool_choice_specific_tool_integration() -> None: + """Test tool_choice parameter with a specific tool using the actual OpenAI API.""" + api_key = os.getenv("OPENAI_API_KEY") + if not api_key: + pytest.skip("OPENAI_API_KEY not found in environment variables") + + def _pass_function(input: str) -> str: + """Simple passthrough function.""" + return f"Processed: {input}" + + def _add_numbers(a: int, b: int) -> int: + """Add two numbers together.""" + return a + b + + model = "gpt-4o-mini" + client = OpenAIChatCompletionClient(model=model, api_key=api_key) + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="_pass_function") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="_add_numbers") + + # Test forcing use of specific tool + result = await client.create( + messages=[UserMessage(content="Process the word 'hello'", source="user")], + tools=[pass_tool, add_tool], + tool_choice=pass_tool, # Force use of specific tool + ) + + assert isinstance(result.content, list) + assert len(result.content) == 1 + assert isinstance(result.content[0], FunctionCall) + assert result.content[0].name == "_pass_function" + assert result.finish_reason == "function_calls" + assert result.usage is not None + + +@pytest.mark.asyncio +async def test_openai_tool_choice_auto_integration() -> None: + """Test tool_choice parameter with 'auto' setting using the actual OpenAI API.""" + api_key = os.getenv("OPENAI_API_KEY") + if not api_key: + pytest.skip("OPENAI_API_KEY not found in environment variables") + + def _pass_function(input: str) -> str: + """Simple passthrough function.""" + return f"Processed: {input}" + + def _add_numbers(a: int, b: int) -> int: + """Add two numbers together.""" + return a + b + + model = "gpt-4o-mini" + client = OpenAIChatCompletionClient(model=model, api_key=api_key) + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="_pass_function") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="_add_numbers") + + # Test auto tool choice - model should choose to use add_numbers for math + result = await client.create( + messages=[UserMessage(content="What is 15 plus 27?", source="user")], + tools=[pass_tool, add_tool], + tool_choice="auto", # Let model choose + ) + + assert isinstance(result.content, list) + assert len(result.content) == 1 + assert isinstance(result.content[0], FunctionCall) + assert result.content[0].name == "_add_numbers" + assert result.finish_reason == "function_calls" + assert result.usage is not None + + # Parse arguments to verify correct values + args = json.loads(result.content[0].arguments) + assert args["a"] == 15 + assert args["b"] == 27 + + +@pytest.mark.asyncio +async def test_openai_tool_choice_none_integration() -> None: + """Test tool_choice parameter with 'none' setting using the actual OpenAI API.""" + api_key = os.getenv("OPENAI_API_KEY") + if not api_key: + pytest.skip("OPENAI_API_KEY not found in environment variables") + + def _pass_function(input: str) -> str: + """Simple passthrough function.""" + return f"Processed: {input}" + + model = "gpt-4o-mini" + client = OpenAIChatCompletionClient(model=model, api_key=api_key) + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="_pass_function") + + # Test none tool choice - model should not use any tools + result = await client.create( + messages=[UserMessage(content="Hello there, how are you?", source="user")], + tools=[pass_tool], + tool_choice="none", # Disable tool usage + ) + + assert isinstance(result.content, str) + assert len(result.content) > 0 + assert result.finish_reason == "stop" + assert result.usage is not None + + +@pytest.mark.asyncio +async def test_openai_tool_choice_required_integration() -> None: + """Test tool_choice parameter with 'required' setting using the actual OpenAI API.""" + api_key = os.getenv("OPENAI_API_KEY") + if not api_key: + pytest.skip("OPENAI_API_KEY not found in environment variables") + + def _pass_function(input: str) -> str: + """Simple passthrough function.""" + return f"Processed: {input}" + + def _add_numbers(a: int, b: int) -> int: + """Add two numbers together.""" + return a + b + + model = "gpt-4o-mini" + client = OpenAIChatCompletionClient(model=model, api_key=api_key) + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="_pass_function") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="_add_numbers") + + # Test required tool choice - model must use a tool even for general conversation + result = await client.create( + messages=[UserMessage(content="Say hello to me", source="user")], + tools=[pass_tool, add_tool], + tool_choice="required", # Force tool usage + ) + + assert isinstance(result.content, list) + assert len(result.content) == 1 + assert isinstance(result.content[0], FunctionCall) + assert result.content[0].name in ["_pass_function", "_add_numbers"] + assert result.finish_reason == "function_calls" + assert result.usage is not None + + +@pytest.mark.asyncio +async def test_openai_tool_choice_validation_error_integration() -> None: + """Test tool_choice validation with invalid tool reference using the actual OpenAI API.""" + api_key = os.getenv("OPENAI_API_KEY") + if not api_key: + pytest.skip("OPENAI_API_KEY not found in environment variables") + + def _pass_function(input: str) -> str: + """Simple passthrough function.""" + return f"Processed: {input}" + + def _add_numbers(a: int, b: int) -> int: + """Add two numbers together.""" + return a + b + + def _different_function(text: str) -> str: + """Different function.""" + return text + + model = "gpt-4o-mini" + client = OpenAIChatCompletionClient(model=model, api_key=api_key) + + # Define tools + pass_tool = FunctionTool(_pass_function, description="Process input text", name="_pass_function") + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="_add_numbers") + different_tool = FunctionTool(_different_function, description="Different tool", name="_different_function") + + messages = [UserMessage(content="Hello there", source="user")] + + # Test with a tool that's not in the tools list + with pytest.raises( + ValueError, match="tool_choice references '_different_function' but it's not in the provided tools" + ): + await client.create( + messages=messages, + tools=[pass_tool, add_tool], + tool_choice=different_tool, # This tool is not in the tools list + ) + + # TODO: add integration tests for Azure OpenAI using AAD token. diff --git a/python/packages/autogen-ext/tests/models/test_sk_chat_completion_adapter.py b/python/packages/autogen-ext/tests/models/test_sk_chat_completion_adapter.py index 0f694b8492ac..300ae0982904 100644 --- a/python/packages/autogen-ext/tests/models/test_sk_chat_completion_adapter.py +++ b/python/packages/autogen-ext/tests/models/test_sk_chat_completion_adapter.py @@ -16,7 +16,7 @@ SystemMessage, UserMessage, ) -from autogen_core.tools import BaseTool +from autogen_core.tools import BaseTool, ParametersSchema, ToolSchema from autogen_ext.models.semantic_kernel import SKChatCompletionAdapter from openai.types.chat.chat_completion_chunk import ( ChatCompletionChunk, @@ -335,6 +335,45 @@ async def test_sk_chat_completion_with_tools(sk_client: AzureChatCompletion) -> assert not result.cached +@pytest.mark.asyncio +async def test_sk_chat_completion_with_prompt_tools(sk_client: AzureChatCompletion) -> None: + # Create adapter + adapter = SKChatCompletionAdapter(sk_client) + + # Create kernel + kernel = Kernel(memory=NullMemory()) + + # Create calculator tool instance + tool: ToolSchema = ToolSchema( + name="calculator", + description="Add two numbers together", + parameters=ParametersSchema( + type="object", + properties={ + "a": {"type": "number", "description": "First number"}, + "b": {"type": "number", "description": "Second number"}, + }, + required=["a", "b"], + ), + ) + + # Test messages + messages: list[LLMMessage] = [ + SystemMessage(content="You are a helpful assistant."), + UserMessage(content="What is 2 + 2?", source="user"), + ] + + # Call create with tool + result = await adapter.create(messages=messages, tools=[tool], extra_create_args={"kernel": kernel}) + + # Verify response + assert isinstance(result.content, list) + assert result.finish_reason == "function_calls" + assert result.usage.prompt_tokens >= 0 + assert result.usage.completion_tokens >= 0 + assert not result.cached + + @pytest.mark.asyncio async def test_sk_chat_completion_without_tools( sk_client: AzureChatCompletion, caplog: pytest.LogCaptureFixture diff --git a/python/packages/autogen-ext/tests/task_centric_memory/utils.py b/python/packages/autogen-ext/tests/task_centric_memory/utils.py index 196176915061..9c1a870bef34 100644 --- a/python/packages/autogen-ext/tests/task_centric_memory/utils.py +++ b/python/packages/autogen-ext/tests/task_centric_memory/utils.py @@ -28,4 +28,4 @@ def load_yaml_file(file_path: str) -> Any: Opens a file and returns its contents. """ with open(file_path, "r") as file: - return yaml.load(file, Loader=yaml.FullLoader) + return yaml.safe_load(file) diff --git a/python/packages/autogen-ext/tests/teams/__init__.py b/python/packages/autogen-ext/tests/teams/__init__.py new file mode 100644 index 000000000000..5d2a4d07959f --- /dev/null +++ b/python/packages/autogen-ext/tests/teams/__init__.py @@ -0,0 +1 @@ +"""Init file for teams tests.""" diff --git a/python/packages/autogen-ext/tests/teams/test_magentic_one.py b/python/packages/autogen-ext/tests/teams/test_magentic_one.py new file mode 100644 index 000000000000..4050f5fe4177 --- /dev/null +++ b/python/packages/autogen-ext/tests/teams/test_magentic_one.py @@ -0,0 +1,257 @@ +"""Tests for MagenticOne team.""" + +import os +import warnings +from unittest.mock import Mock, patch + +import pytest +from autogen_agentchat.agents import CodeExecutorAgent +from autogen_agentchat.agents._code_executor_agent import ApprovalRequest, ApprovalResponse +from autogen_core.models import ChatCompletionClient +from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor +from autogen_ext.teams.magentic_one import MagenticOne + + +def docker_tests_enabled() -> bool: + """Check if Docker tests should be enabled.""" + if os.environ.get("SKIP_DOCKER", "unset").lower() == "true": + return False + + try: + import docker + from docker.errors import DockerException + except ImportError: + return False + + try: + client = docker.from_env() + client.ping() # type: ignore + return True + except DockerException: + return False + + +def _is_docker_available() -> bool: + """Local implementation of Docker availability check.""" + return docker_tests_enabled() + + +@pytest.fixture +def mock_chat_client() -> Mock: + """Create a mock chat completion client.""" + mock_client = Mock(spec=ChatCompletionClient) + mock_client.model_info = {"function_calling": True, "json_output": True, "vision": True} + return mock_client + + +def approval_function_allow_all(request: ApprovalRequest) -> ApprovalResponse: + """Test approval function that allows all code execution.""" + return ApprovalResponse(approved=True, reason="Test approval - all code allowed") + + +def approval_function_deny_all(request: ApprovalRequest) -> ApprovalResponse: + """Test approval function that denies all code execution.""" + return ApprovalResponse(approved=False, reason="Test approval - all code denied") + + +@pytest.mark.skipif(not docker_tests_enabled(), reason="Docker is not available") +def test_magentic_one_uses_docker_by_default(mock_chat_client: Mock) -> None: + """Test that MagenticOne uses Docker code executor by default when Docker is available.""" + from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + + m1 = MagenticOne(client=mock_chat_client) + + # Find the CodeExecutorAgent in the participants list + code_executor_agent = None + for agent in m1._participants: # type: ignore[reportPrivateUsage] + if isinstance(agent, CodeExecutorAgent): + code_executor_agent = agent + break + + assert code_executor_agent is not None, "CodeExecutorAgent not found" + assert isinstance( + code_executor_agent._code_executor, # type: ignore[reportPrivateUsage] + DockerCommandLineCodeExecutor, # type: ignore[reportPrivateUsage] + ), f"Expected DockerCommandLineCodeExecutor, got {type(code_executor_agent._code_executor)}" # type: ignore[reportPrivateUsage] + + # Test that no approval function is set by default + assert code_executor_agent._approval_func is None, "Expected no approval function by default" # type: ignore[reportPrivateUsage] + + +@pytest.mark.skipif(not docker_tests_enabled(), reason="Docker is not available") +def test_magentic_one_uses_docker_with_approval_function(mock_chat_client: Mock) -> None: + """Test that MagenticOne uses Docker code executor with approval function when provided.""" + from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + + m1 = MagenticOne(client=mock_chat_client, approval_func=approval_function_allow_all) + + # Find the CodeExecutorAgent in the participants list + code_executor_agent = None + for agent in m1._participants: # type: ignore[reportPrivateUsage] + if isinstance(agent, CodeExecutorAgent): + code_executor_agent = agent + break + + assert code_executor_agent is not None, "CodeExecutorAgent not found" + assert isinstance( + code_executor_agent._code_executor, # type: ignore[reportPrivateUsage] + DockerCommandLineCodeExecutor, # type: ignore[reportPrivateUsage] + ), f"Expected DockerCommandLineCodeExecutor, got {type(code_executor_agent._code_executor)}" # type: ignore[reportPrivateUsage] + + # Test that approval function is set correctly + assert code_executor_agent._approval_func is approval_function_allow_all, "Expected approval function to be set" # type: ignore[reportPrivateUsage] + + +def test_docker_availability_check() -> None: + """Test the Docker availability check function.""" + # This test should pass regardless of Docker availability + result = _is_docker_available() + assert isinstance(result, bool) + + +@patch("autogen_ext.teams.magentic_one._is_docker_available") +def test_magentic_one_falls_back_to_local_when_docker_unavailable( + mock_docker_check: Mock, mock_chat_client: Mock +) -> None: + """Test that MagenticOne falls back to local executor when Docker is not available.""" + mock_docker_check.return_value = False + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + m1 = MagenticOne(client=mock_chat_client) + + # Find the CodeExecutorAgent in the participants list + code_executor_agent = None + for agent in m1._participants: # type: ignore[reportPrivateUsage] + if isinstance(agent, CodeExecutorAgent): + code_executor_agent = agent + break + + assert code_executor_agent is not None, "CodeExecutorAgent not found" + assert isinstance( + code_executor_agent._code_executor, # type: ignore[reportPrivateUsage] + LocalCommandLineCodeExecutor, # type: ignore[reportPrivateUsage] + ), f"Expected LocalCommandLineCodeExecutor, got {type(code_executor_agent._code_executor)}" # type: ignore[reportPrivateUsage] + + # Test that no approval function is set by default + assert code_executor_agent._approval_func is None, "Expected no approval function by default" # type: ignore[reportPrivateUsage] + + # Check that appropriate warnings were issued + warning_messages = [str(warning.message) for warning in w] + docker_warning_found = any("Docker is not available" in msg for msg in warning_messages) + deprecated_warning_found = any( + "Instantiating MagenticOne without a code_executor is deprecated" in msg for msg in warning_messages + ) + + assert docker_warning_found, f"Docker unavailable warning not found in: {warning_messages}" + assert deprecated_warning_found, f"Deprecation warning not found in: {warning_messages}" + + +@patch("autogen_ext.teams.magentic_one._is_docker_available") +def test_magentic_one_falls_back_to_local_with_approval_function( + mock_docker_check: Mock, mock_chat_client: Mock +) -> None: + """Test that MagenticOne falls back to local executor with approval function when Docker is not available.""" + mock_docker_check.return_value = False + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + m1 = MagenticOne(client=mock_chat_client, approval_func=approval_function_deny_all) + + # Find the CodeExecutorAgent in the participants list + code_executor_agent = None + for agent in m1._participants: # type: ignore[reportPrivateUsage] + if isinstance(agent, CodeExecutorAgent): + code_executor_agent = agent + break + + assert code_executor_agent is not None, "CodeExecutorAgent not found" + assert isinstance( + code_executor_agent._code_executor, # type: ignore[reportPrivateUsage] + LocalCommandLineCodeExecutor, # type: ignore[reportPrivateUsage] + ), f"Expected LocalCommandLineCodeExecutor, got {type(code_executor_agent._code_executor)}" # type: ignore[reportPrivateUsage] + + # Test that approval function is set correctly + assert code_executor_agent._approval_func is approval_function_deny_all, "Expected approval function to be set" # type: ignore[reportPrivateUsage] + + # Check that appropriate warnings were issued + warning_messages = [str(warning.message) for warning in w] + docker_warning_found = any("Docker is not available" in msg for msg in warning_messages) + deprecated_warning_found = any( + "Instantiating MagenticOne without a code_executor is deprecated" in msg for msg in warning_messages + ) + + assert docker_warning_found, f"Docker unavailable warning not found in: {warning_messages}" + assert deprecated_warning_found, f"Deprecation warning not found in: {warning_messages}" + + +def test_magentic_one_with_explicit_code_executor(mock_chat_client: Mock) -> None: + """Test that MagenticOne uses the provided code executor when explicitly given.""" + explicit_executor = LocalCommandLineCodeExecutor() + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + m1 = MagenticOne(client=mock_chat_client, code_executor=explicit_executor) + + # Find the CodeExecutorAgent in the participants list + code_executor_agent = None + for agent in m1._participants: # type: ignore[reportPrivateUsage] + if isinstance(agent, CodeExecutorAgent): + code_executor_agent = agent + break + + assert code_executor_agent is not None, "CodeExecutorAgent not found" + assert code_executor_agent._code_executor is explicit_executor, "Expected the explicitly provided code executor" # type: ignore[reportPrivateUsage] + + # Test that no approval function is set by default + assert code_executor_agent._approval_func is None, "Expected no approval function by default" # type: ignore[reportPrivateUsage] + + # No deprecation warning should be issued when explicitly providing a code executor + warning_messages = [str(warning.message) for warning in w] + deprecated_warning_found = any( + "Instantiating MagenticOne without a code_executor is deprecated" in msg for msg in warning_messages + ) + + assert not deprecated_warning_found, f"Unexpected deprecation warning found: {warning_messages}" + + +def test_magentic_one_with_explicit_code_executor_and_approval_function(mock_chat_client: Mock) -> None: + """Test that MagenticOne uses the provided code executor and approval function when explicitly given.""" + explicit_executor = LocalCommandLineCodeExecutor() + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + m1 = MagenticOne( + client=mock_chat_client, code_executor=explicit_executor, approval_func=approval_function_allow_all + ) + + # Find the CodeExecutorAgent in the participants list + code_executor_agent = None + for agent in m1._participants: # type: ignore[reportPrivateUsage] + if isinstance(agent, CodeExecutorAgent): + code_executor_agent = agent + break + + assert code_executor_agent is not None, "CodeExecutorAgent not found" + assert code_executor_agent._code_executor is explicit_executor, "Expected the explicitly provided code executor" # type: ignore[reportPrivateUsage] + + # Test that approval function is set correctly + assert code_executor_agent._approval_func is approval_function_allow_all, "Expected approval function to be set" # type: ignore[reportPrivateUsage] + + # No deprecation warning should be issued when explicitly providing a code executor + warning_messages = [str(warning.message) for warning in w] + deprecated_warning_found = any( + "Instantiating MagenticOne without a code_executor is deprecated" in msg for msg in warning_messages + ) + + assert not deprecated_warning_found, f"Unexpected deprecation warning found: {warning_messages}" diff --git a/python/packages/autogen-ext/tests/test_azure_ai_agent.py b/python/packages/autogen-ext/tests/test_azure_ai_agent.py new file mode 100644 index 000000000000..a20a49c1d458 --- /dev/null +++ b/python/packages/autogen-ext/tests/test_azure_ai_agent.py @@ -0,0 +1,796 @@ +import json +from asyncio import CancelledError +from types import SimpleNamespace +from typing import Any, AsyncGenerator, List, Optional, Union +from unittest.mock import AsyncMock, MagicMock, call + +import pytest +from autogen_agentchat.base._chat_agent import Response +from autogen_agentchat.messages import TextMessage, ToolCallExecutionEvent +from autogen_core._cancellation_token import CancellationToken +from autogen_core.tools._function_tool import FunctionTool +from autogen_ext.agents.azure._azure_ai_agent import AzureAIAgent +from autogen_ext.agents.azure._types import ListToolType +from azure.ai.agents.models import ( + AzureAISearchToolDefinition, + AzureFunctionToolDefinition, + BingGroundingToolDefinition, + CodeInterpreterToolDefinition, + FilePurpose, + FileSearchToolDefinition, + FileState, + RequiredAction, + RunStatus, + SubmitToolOutputsAction, + ThreadMessage, +) +from azure.ai.projects.aio import AIProjectClient + + +class FakeText: + def __init__(self, value: str) -> None: + self.value = value + + +class FakeTextContent: + def __init__(self, text: str) -> None: + self.type = "text" + self.text = FakeText(text) + + +class FakeMessage: + def __init__(self, id: str, text: str) -> None: + self.id = id + # The agent expects content to be a list of objects with a "type" attribute. + self.content = [FakeTextContent(text)] + self.role = "user" + + @property + def text_messages(self) -> List[FakeTextContent]: + """Returns all text message contents in the messages. + + :rtype: List[FakeTextContent] + """ + if not self.content: + return [] + return [content for content in self.content if isinstance(content, FakeTextContent)] + + +class FakeMessageUrlCitationDetails: + def __init__(self, url: str, title: str) -> None: + self.url = url + self.title = title + + +class FakeTextUrlCitationAnnotation: + def __init__(self, citation_details: FakeMessageUrlCitationDetails, text: str) -> None: + self.type = "url_citation" + self.url_citation = citation_details + self.text = text + + +class FakeTextFileCitationDetails: + def __init__(self, file_id: str, quote: str) -> None: + self.file_id = file_id + self.quote = quote + + +class FakeTextFileCitationAnnotation: + def __init__(self, citation_details: FakeTextFileCitationDetails) -> None: + self.type = "file_citation" + self.file_citation = citation_details + + +class FakeMessageWithUrlCitationAnnotation: + def __init__(self, id: str, text: str, annotations: list[FakeTextUrlCitationAnnotation]) -> None: + self.id = id + # The agent expects content to be a list of objects with a "type" attribute. + self.content = [FakeTextContent(text)] + self.role = "user" + self._annotations = annotations + + @property + def text_messages(self) -> List[FakeTextContent]: + """Returns all text message contents in the messages. + + :rtype: List[FakeTextContent] + """ + if not self.content: + return [] + return [content for content in self.content if isinstance(content, FakeTextContent)] + + @property + def url_citation_annotations(self) -> List[FakeTextUrlCitationAnnotation]: + """Returns all URL citation annotations from text message annotations in the messages. + + :rtype: List[FakeTextUrlCitationAnnotation] + """ + return self._annotations + + +class FakeMessageWithFileCitationAnnotation: + def __init__(self, id: str, text: str, annotations: list[FakeTextFileCitationAnnotation]) -> None: + self.id = id + # The agent expects content to be a list of objects with a "type" attribute. + self.content = [FakeTextContent(text)] + self.role = "user" + self._annotations = annotations + + @property + def text_messages(self) -> List[FakeTextContent]: + """Returns all text message contents in the messages. + + :rtype: List[FakeTextContent] + """ + if not self.content: + return [] + return [content for content in self.content if isinstance(content, FakeTextContent)] + + @property + def file_citation_annotations(self) -> List[FakeTextFileCitationAnnotation]: + """Returns all URL citation annotations from text message annotations in the messages. + + :rtype: List[FakeTextFileCitationAnnotation] + """ + return self._annotations + + +class FakeMessageWithAnnotation: + def __init__(self, id: str, text: str, annotations: list[FakeTextUrlCitationAnnotation]) -> None: + self.id = id + # The agent expects content to be a list of objects with a "type" attribute. + self.content = [FakeTextContent(text)] + self.role = "user" + self.annotations = annotations + + @property + def text_messages(self) -> List[FakeTextContent]: + """Returns all text message contents in the messages. + + :rtype: List[FakeTextContent] + """ + if not self.content: + return [] + return [content for content in self.content if isinstance(content, FakeTextContent)] + + +FakeMessageType = Union[ + ThreadMessage + | FakeMessage + | FakeMessageWithAnnotation + | FakeMessageWithUrlCitationAnnotation + | FakeMessageWithFileCitationAnnotation +] + + +async def mock_messages_list(**kwargs: Any) -> AsyncGenerator[FakeMessage, None]: + """Mock async generator for messages.list()""" + messages = [FakeMessage("msg-mock", "response")] + for message in messages: + yield message + + +async def mock_messages_list_empty(**kwargs: Any) -> AsyncGenerator[FakeMessage, None]: + """Mock async generator that yields no messages""" + # This generator yields nothing, simulating an empty message list + return + yield # This line is never reached but makes this a generator + + +async def mock_messages_list_multiple(**kwargs: Any) -> AsyncGenerator[FakeMessage, None]: + """Mock async generator for multiple messages (pagination test)""" + messages = [ + FakeMessage("msg-mock-1", "response-1"), + FakeMessage("msg-mock-2", "response-2"), + ] + for message in messages: + yield message + + +def create_agent( + mock_project_client: MagicMock, + tools: Optional[ListToolType] = None, + agent_name: str = "test_agent", + description: str = "Test Azure AI Agent", + instructions: str = "Test instructions", + agent_id: Optional[str] = None, + thread_id: Optional[str] = None, +) -> AzureAIAgent: + return AzureAIAgent( + name=agent_name, + description=description, + project_client=mock_project_client, + deployment_name="test_model", + tools=tools, + instructions=instructions, + agent_id=agent_id, + thread_id=thread_id, + ) + + +@pytest.fixture +def mock_project_client() -> MagicMock: + client = MagicMock(spec=AIProjectClient) + + # Create separate operation groups to match the actual SDK structure + client.agents = MagicMock() + client.runs = MagicMock() + client.messages = MagicMock() + client.threads = MagicMock() + client.files = MagicMock() + client.vector_stores = MagicMock() + client.vector_store_files = MagicMock() + client.vector_store_file_batches = MagicMock() + + # Agent operations + client.agents.create_agent = AsyncMock(return_value=MagicMock(id="assistant-mock")) + client.agents.get_agent = AsyncMock(return_value=MagicMock(id="assistant-mock")) + client.agents.update_agent = AsyncMock() + client.agents.delete_agent = AsyncMock() + + agent_run = MagicMock() + agent_run.id = "run-mock" + agent_run.status = RunStatus.COMPLETED + + client.agents.runs = MagicMock() + client.agents.runs.create = AsyncMock(return_value=agent_run) + client.agents.runs.get = AsyncMock(return_value=agent_run) + client.agents.runs.submit_tool_outputs = AsyncMock(return_value=agent_run) + + client.agents.messages = MagicMock() + client.agents.messages.list = mock_messages_list + client.agents.messages.create = AsyncMock() + + client.agents.threads = MagicMock() + client.agents.threads.get = AsyncMock(return_value=MagicMock(id="thread-mock")) + client.agents.threads.create = AsyncMock(return_value=MagicMock(id="thread-mock")) + client.agents.threads.update = AsyncMock() + + client.agents.files = MagicMock() + client.agents.files.upload_and_poll = AsyncMock(return_value=MagicMock(id="file-mock", status=FileState.PROCESSED)) + + client.agents.vector_stores = MagicMock() + client.agents.vector_stores.create_and_poll = AsyncMock(return_value=MagicMock(id="vector_store_id")) + client.agents.vector_store_file_batches = MagicMock() + client.agents.vector_store_file_batches.create_and_poll = AsyncMock() + + return client + + +@pytest.mark.asyncio +async def test_azure_ai_agent_initialization(mock_project_client: MagicMock) -> None: + agent = create_agent(mock_project_client, ["file_search"]) + + assert agent.name == "test_agent" + assert agent.description == "Test Azure AI Agent" + assert agent.deployment_name == "test_model" + assert agent.instructions == "Test instructions" + assert len(agent.tools) == 1 + + +@pytest.mark.asyncio +async def test_on_messages(mock_project_client: MagicMock) -> None: + agent = create_agent(mock_project_client) + + messages = [TextMessage(content="Hello", source="user")] + response = await agent.on_messages(messages) + + assert response is not None + + +@pytest.mark.asyncio +async def test_on_reset(mock_project_client: MagicMock) -> None: + agent = create_agent(mock_project_client) + + await agent.on_reset(CancellationToken()) + + # The agent might call create_thread multiple times during initialization, so check if it was called at least once + assert mock_project_client.agents.threads.create.call_count > 0 + + +@pytest.mark.asyncio +async def test_save_and_load_state(mock_project_client: MagicMock) -> None: + agent = create_agent(mock_project_client, agent_id="agent-mock", thread_id="thread-mock") + + state = await agent.save_state() + assert state is not None + + await agent.load_state(state) + + assert agent.agent_id == state["agent_id"] + # assert agent._init_thread_id == state["thread_id"] + + +@pytest.mark.asyncio +async def test_on_upload_for_code_interpreter(mock_project_client: MagicMock) -> None: + file_mock = MagicMock() + file_mock.id = "file-mock" + file_mock.status = FileState.PROCESSED + + thread_mock = MagicMock() + thread_mock.tool_resources = MagicMock() + thread_mock.tool_resources.code_interpreter = MagicMock() + thread_mock.tool_resources.code_interpreter.file_ids = [] # Set as a valid list + + mock_project_client.agents.files.upload_and_poll = AsyncMock(return_value=file_mock) + mock_project_client.agents.threads.get = AsyncMock(return_value=thread_mock) + mock_project_client.agents.threads.update = AsyncMock() + + agent = create_agent( + mock_project_client, + ) + + file_paths = ["test_file_1.txt", "test_file_2.txt"] + await agent.on_upload_for_code_interpreter(file_paths) + + mock_project_client.agents.files.upload_and_poll.assert_called() + mock_project_client.agents.threads.get.assert_called_once() + mock_project_client.agents.threads.update.assert_called_once() + + +@pytest.mark.asyncio +async def test_on_upload_for_file_search(mock_project_client: MagicMock) -> None: + file_mock = MagicMock() + file_mock.id = "file-mock" + file_mock.status = FileState.PROCESSED # Set a valid status + + mock_project_client.agents.files.upload_and_poll = AsyncMock(return_value=file_mock) + mock_project_client.agents.vector_stores.create_and_poll = AsyncMock(return_value=MagicMock(id="vector_store_id")) + mock_project_client.agents.update_agent = AsyncMock() + mock_project_client.agents.vector_store_file_batches.create_and_poll = AsyncMock() + + agent = create_agent(mock_project_client, tools=["file_search"]) + + file_paths = ["test_file_1.txt", "test_file_2.txt"] + await agent.on_upload_for_file_search(file_paths, cancellation_token=CancellationToken()) + + mock_project_client.agents.files.upload_and_poll.assert_called() + mock_project_client.agents.vector_stores.create_and_poll.assert_called_once() + mock_project_client.agents.update_agent.assert_called_once() + mock_project_client.agents.vector_store_file_batches.create_and_poll.assert_called_once() + + +@pytest.mark.asyncio +async def test_upload_files(mock_project_client: MagicMock) -> None: + mock_project_client.agents.vector_store_file_batches.create_and_poll = AsyncMock() + + mock_project_client.agents.update_agent = AsyncMock() + mock_project_client.agents.vector_stores.create_and_poll = AsyncMock(return_value=MagicMock(id="vector_store_id")) + + mock_project_client.agents.files.upload_and_poll = AsyncMock( + return_value=MagicMock(id="file-id", status=FileState.PROCESSED) + ) + + agent = create_agent(mock_project_client, tools=["file_search"]) + + await agent.on_upload_for_file_search(["test_file.txt"], cancellation_token=CancellationToken()) + + mock_project_client.agents.files.upload_and_poll.assert_any_await( + file_path="test_file.txt", purpose=FilePurpose.AGENTS, polling_interval=0.5 + ) + + +@pytest.mark.asyncio +async def test_on_messages_stream(mock_project_client: MagicMock) -> None: + mock_project_client.agents.runs.create = AsyncMock( # Corrected path + return_value=MagicMock(id="run-id", status=RunStatus.COMPLETED) + ) + mock_project_client.agents.messages.list = mock_messages_list # Corrected path + + agent = create_agent(mock_project_client) + + messages = [TextMessage(content="Hello", source="user")] + async for response in agent.on_messages_stream(messages): + assert isinstance(response, Response) + assert response.chat_message.to_model_message().content == "response" + + +@pytest.mark.asyncio +async def test_on_messages_stream_with_tool(mock_project_client: MagicMock) -> None: + agent = create_agent(mock_project_client, tools=["file_search"]) + + messages = [TextMessage(content="Hello", source="user")] + async for response in agent.on_messages_stream(messages): + assert isinstance(response, Response) + assert response.chat_message.to_model_message().content == "response" + + +@pytest.mark.asyncio +async def test_thread_id_validation(mock_project_client: MagicMock) -> None: + agent = create_agent(mock_project_client) + + with pytest.raises(ValueError, match="Thread not"): + _ = agent.thread_id # Using _ for intentionally unused variable + + +@pytest.mark.asyncio +async def test_get_agent_id_validation(mock_project_client: MagicMock) -> None: + agent = create_agent(mock_project_client) + + with pytest.raises(ValueError, match="Agent not"): + _ = agent.agent_id # Using _ for intentionally unused variable + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_name, should_raise_error", + [ + ("file_search", False), + ("code_interpreter", False), + ("bing_grounding", False), + ("azure_function", False), + ("azure_ai_search", False), + # ("sharepoint_grounding", False), + ("unknown_tool", True), + ], +) +async def test_adding_tools_as_literals( + mock_project_client: MagicMock, tool_name: Any, should_raise_error: bool +) -> None: + if should_raise_error: + with pytest.raises(ValueError, match=tool_name): + agent = create_agent(mock_project_client, tools=[tool_name]) # mypy ignore + else: + agent = create_agent(mock_project_client, tools=[tool_name]) + assert agent.tools[0].type == tool_name + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_definition", + [ + FileSearchToolDefinition(), + CodeInterpreterToolDefinition(), + BingGroundingToolDefinition(), # type: ignore + AzureFunctionToolDefinition(), # type: ignore + AzureAISearchToolDefinition(), + # SharepointToolDefinition(), # type: ignore + ], +) +async def test_adding_tools_as_typed_definition(mock_project_client: MagicMock, tool_definition: Any) -> None: + agent = create_agent(mock_project_client, tools=[tool_definition]) + + assert len(agent.tools) == 1 + assert agent.tools[0].type == tool_definition.type + + +@pytest.mark.asyncio +async def test_adding_callable_func_as_tool(mock_project_client: MagicMock) -> None: + def mock_tool_func() -> None: + """Mock tool function.""" + pass + + agent = create_agent(mock_project_client, tools=[mock_tool_func]) + assert len(agent.tools) == 1 + + assert agent.tools[0].type == "function" + + +@pytest.mark.asyncio +async def test_adding_core_autogen_tool(mock_project_client: MagicMock) -> None: + def mock_tool_func() -> None: + """Mock tool function.""" + pass + + tool = FunctionTool( + func=mock_tool_func, + name="mock_tool", + description="Mock tool function", + ) + + agent = create_agent(mock_project_client, tools=[tool]) + + assert len(agent.tools) == 1 + assert agent.tools[0].type == "function" + + +@pytest.mark.asyncio +async def test_adding_core_autogen_tool_without_doc_string(mock_project_client: MagicMock) -> None: + def mock_tool_func() -> None: + pass + + agent = create_agent(mock_project_client, tools=[mock_tool_func]) + + assert len(agent.tools) == 1 + assert agent.tools[0].type == "function" + assert agent.tools[0].function.description == "" # type: ignore + + +@pytest.mark.asyncio +async def test_adding_unsupported_tool(mock_project_client: MagicMock) -> None: + tool_name: Any = 5 + + with pytest.raises(ValueError, match="class 'int'"): + create_agent(mock_project_client, tools=[tool_name]) + + +@pytest.mark.asyncio +async def test_agent_initialization_with_no_agent_id(mock_project_client: MagicMock) -> None: + agent = create_agent(mock_project_client) + + await agent.on_messages([TextMessage(content="Hello", source="user")]) + + mock_project_client.agents.create_agent.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_agent_initialization_with_agent_id(mock_project_client: MagicMock) -> None: + agent = create_agent(mock_project_client, agent_id="agent-mock") + + await agent.on_messages([TextMessage(content="Hello", source="user")]) + + mock_project_client.agents.get_agent.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_agent_initialization_with_no_thread_id(mock_project_client: MagicMock) -> None: + agent = create_agent(mock_project_client) + + await agent.on_messages([TextMessage(content="Hello", source="user")]) + + mock_project_client.agents.threads.create.assert_awaited_once() # Corrected path + + +@pytest.mark.asyncio +async def test_agent_initialization_with_thread_id(mock_project_client: MagicMock) -> None: + agent = create_agent(mock_project_client, thread_id="thread-mock") + + await agent.on_messages([TextMessage(content="Hello", source="user")]) + + mock_project_client.agents.threads.get.assert_awaited_once() # Corrected path + + +@pytest.mark.asyncio +async def test_agent_initialization_fetching_multiple_pages_of_thread_messages(mock_project_client: MagicMock) -> None: + mock_project_client.agents.threads.get = AsyncMock(return_value=MagicMock(id="thread-id")) # Corrected path + # Mock the list_messages method to return multiple messages + mock_project_client.agents.messages.list = mock_messages_list_multiple # Corrected path + + agent = create_agent(mock_project_client, thread_id="thread-id") + + def assert_messages(actual: list[str], expected: List[str]) -> None: + assert len(actual) == len(expected) + for i in range(len(actual)): + assert actual[i] in expected + + try: + await agent.on_messages([TextMessage(content="Hello", source="user")]) + + state = await agent.save_state() + assert state is not None + assert len(state["initial_message_ids"]) == 2 + assert_messages(state["initial_message_ids"], ["msg-mock-1", "msg-mock-2"]) + except StopAsyncIteration: + # Handle the StopAsyncIteration exception to allow the test to continue + pass + + +@pytest.mark.asyncio +async def test_on_messages_with_cancellation(mock_project_client: MagicMock) -> None: + agent = create_agent(mock_project_client) + + # Create a cancellation token that's already cancelled + token = CancellationToken() + token.cancel() + + messages = [TextMessage(content="Hello", source="user")] + + with pytest.raises(CancelledError): + await agent.on_messages(messages, token) + + +def mock_run(action: str, run_id: str, required_action: Optional[RequiredAction] = None) -> MagicMock: + run = MagicMock() + run.id = run_id + run.status = action + run.required_action = required_action + return run + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_name, registered_tools, error", + [ + ( + "function", + [ + FunctionTool( + func=lambda: None, + name="mock_tool", + description="Mock tool function", + ) + ], + "is not available", + ), + ("function", None, "No tools"), + ], +) +async def test_on_messages_return_required_action_with_no_tool_raise_error( + mock_project_client: MagicMock, tool_name: str, registered_tools: ListToolType, error: str +) -> None: + agent = create_agent(mock_project_client, tools=registered_tools) + + complete_run = mock_run(RunStatus.COMPLETED, "run-mock") + mock_project_client.agents.runs.submit_tool_outputs = AsyncMock(return_value=complete_run) # Corrected path + + required_action = SubmitToolOutputsAction( + submit_tool_outputs=SimpleNamespace( # type: ignore + tool_calls=[ + SimpleNamespace( + type="function", + id="tool-mock", + name=tool_name, + function=SimpleNamespace(arguments={}, name="function"), + ) + ] + ) + ) + + required_action.submit_tool_outputs = SimpleNamespace( # type: ignore + tool_calls=[ + SimpleNamespace( + type="function", id="tool-mock", name=tool_name, function=SimpleNamespace(arguments={}, name="function") + ) + ] + ) # mypy ignore + + requires_action_run = mock_run(RunStatus.REQUIRES_ACTION, "run-mock", required_action) + mock_project_client.agents.runs.get = AsyncMock(side_effect=[requires_action_run, complete_run]) # Corrected path + + messages = [TextMessage(content="Hello", source="user")] + + response: Response = await agent.on_messages(messages) + + # check why there are 2 inner messages + tool_call_events = [event for event in response.inner_messages if isinstance(event, ToolCallExecutionEvent)] # type: ignore + assert len(tool_call_events) == 1 + + event: ToolCallExecutionEvent = tool_call_events[0] + assert event.content[0].is_error is True + assert event.content[0].content.find(error) != -1 + + +@pytest.mark.asyncio +async def test_on_message_raise_error_when_stream_return_nothing(mock_project_client: MagicMock) -> None: + agent = create_agent(mock_project_client) + + messages = [TextMessage(content="Hello", source="user")] + agent.on_messages_stream = MagicMock(name="on_messages_stream") # type: ignore + agent.on_messages_stream.__aiter__.return_value = [] + + with pytest.raises(AssertionError, match="have returned the final result"): + await agent.on_messages(messages) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "file_paths, file_status, should_raise_error", + [ + (["file1.txt", "file2.txt"], FileState.PROCESSED, False), + (["file3.txt"], FileState.ERROR, True), + ], +) +async def test_uploading_multiple_files( + mock_project_client: MagicMock, file_paths: list[str], file_status: FileState, should_raise_error: bool +) -> None: + agent = create_agent(mock_project_client) + + file_mock = MagicMock(id="file-id", status=file_status) + mock_project_client.agents.threads.update = AsyncMock() + mock_project_client.agents.files.upload_and_poll = AsyncMock(return_value=file_mock) + + async def upload_files() -> None: + await agent.on_upload_for_code_interpreter( + file_paths, + cancellation_token=CancellationToken(), + polling_interval=0.1, + ) + + if should_raise_error: + with pytest.raises(ValueError, match="upload failed with status"): # Changed from Exception to ValueError + await upload_files() + else: + await upload_files() + + mock_project_client.agents.files.upload_and_poll.assert_has_calls( + [call(file_path=file_path, purpose=FilePurpose.AGENTS, polling_interval=0.1) for file_path in file_paths] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fake_message, url, title", + [ + ( + FakeMessageWithAnnotation( + "msg-mock-1", + "response-1", + [FakeTextUrlCitationAnnotation(FakeMessageUrlCitationDetails("url1", "title1"), "text")], + ), + "url1", + "title1", + ), + ( + FakeMessageWithUrlCitationAnnotation( + "msg-mock-2", + "response-2", + [FakeTextUrlCitationAnnotation(FakeMessageUrlCitationDetails("url2", "title2"), "text")], + ), + "url2", + "title2", + ), + ], +) +async def test_on_message_stream_mapping_url_citation( + mock_project_client: MagicMock, + fake_message: FakeMessageWithAnnotation | FakeMessageWithUrlCitationAnnotation, + url: str, + title: str, +) -> None: + mock_project_client.agents.runs.create = AsyncMock( # Corrected path and method name + return_value=MagicMock(id="run-id", status=RunStatus.COMPLETED) + ) + + async def mock_messages_list_with_citation( + **kwargs: Any, + ) -> AsyncGenerator[FakeMessageWithAnnotation | FakeMessageWithUrlCitationAnnotation, None]: + """Mock async generator for messages with citation""" + yield fake_message + + mock_project_client.agents.messages.list = mock_messages_list_with_citation + + agent = create_agent(mock_project_client) + + messages = [TextMessage(content="Hello", source="user")] + + async for response in agent.on_messages_stream(messages): + assert isinstance(response, Response) + assert response.chat_message is not None + assert response.chat_message.metadata is not None + + citations = json.loads(response.chat_message.metadata["citations"]) + assert citations is not None + + assert len(citations) == 1 + + assert citations[0]["url"] == url + assert citations[0]["title"] == title + + +@pytest.mark.asyncio +async def test_on_message_stream_mapping_file_citation(mock_project_client: MagicMock) -> None: + mock_project_client.agents.create_run = AsyncMock(return_value=MagicMock(id="run-id", status=RunStatus.COMPLETED)) + + expected_file_id = "file_id_1" + expected_quote = "this part of a file" + + fake_message = FakeMessageWithFileCitationAnnotation( + "msg-mock-1", + "response-1", + [FakeTextFileCitationAnnotation(FakeTextFileCitationDetails(expected_file_id, expected_quote))], + ) + + async def mock_messages_list_with_file_citation( + **kwargs: Any, + ) -> AsyncGenerator[FakeMessageWithFileCitationAnnotation, None]: + """Mock async generator for messages with file citation""" + yield fake_message + + mock_project_client.agents.messages.list = mock_messages_list_with_file_citation + + agent = create_agent(mock_project_client) + + messages = [TextMessage(content="Hello", source="user")] + + async for response in agent.on_messages_stream(messages): + assert isinstance(response, Response) + assert response.chat_message is not None + assert response.chat_message.metadata is not None + + citations = json.loads(response.chat_message.metadata["citations"]) + assert citations is not None + + assert len(citations) == 1 + + assert citations[0]["file_id"] == expected_file_id + assert citations[0]["text"] == expected_quote diff --git a/python/packages/autogen-ext/tests/test_filesurfer_agent.py b/python/packages/autogen-ext/tests/test_filesurfer_agent.py index 470bb270a9ef..de2bbfec837b 100644 --- a/python/packages/autogen-ext/tests/test_filesurfer_agent.py +++ b/python/packages/autogen-ext/tests/test_filesurfer_agent.py @@ -8,6 +8,7 @@ import aiofiles import pytest from autogen_agentchat import EVENT_LOGGER_NAME +from autogen_agentchat.messages import TextMessage from autogen_ext.agents.file_surfer import FileSurfer from autogen_ext.models.openai import OpenAIChatCompletionClient from openai.resources.chat.completions import AsyncCompletions @@ -31,7 +32,7 @@ def emit(self, record: logging.LogRecord) -> None: record.msg = json.dumps( { "timestamp": ts, - "message": record.msg.model_dump(), + "message": record.msg.model_dump_json(indent=2), "type": record.msg.__class__.__name__, }, ) @@ -73,7 +74,7 @@ async def test_run_filesurfer(monkeypatch: pytest.MonkeyPatch) -> None: """) # Mock the API calls - model = "gpt-4o-2024-05-13" + model = "gpt-4.1-nano-2025-04-14" chat_completions = [ ChatCompletion( id="id1", @@ -140,9 +141,11 @@ async def test_run_filesurfer(monkeypatch: pytest.MonkeyPatch) -> None: # Get the FileSurfer to read the file, and the directory assert agent._name == "FileSurfer" # pyright: ignore[reportPrivateUsage] result = await agent.run(task="Please read the test file") + assert isinstance(result.messages[1], TextMessage) assert "# FileSurfer test H1" in result.messages[1].content result = await agent.run(task="Please read the test directory") + assert isinstance(result.messages[1], TextMessage) assert "# Index of " in result.messages[1].content assert "test_filesurfer_agent.html" in result.messages[1].content @@ -150,7 +153,7 @@ async def test_run_filesurfer(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.mark.asyncio async def test_file_surfer_serialization() -> None: """Test that FileSurfer can be serialized and deserialized properly.""" - model = "gpt-4o-2024-05-13" + model = "gpt-4.1-nano-2025-04-14" agent = FileSurfer( "FileSurfer", model_client=OpenAIChatCompletionClient(model=model, api_key=""), diff --git a/python/packages/autogen-ext/tests/test_openai_agent.py b/python/packages/autogen-ext/tests/test_openai_agent.py new file mode 100644 index 000000000000..66fd998de88e --- /dev/null +++ b/python/packages/autogen-ext/tests/test_openai_agent.py @@ -0,0 +1,513 @@ +import json +from typing import Any, AsyncGenerator, Dict, List, Mapping, Optional, Type, Union, cast +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from autogen_agentchat.base import Response +from autogen_agentchat.messages import BaseChatMessage, MultiModalMessage, TextMessage +from autogen_core import CancellationToken, Image +from autogen_core.models import UserMessage +from autogen_core.tools import Tool, ToolSchema +from autogen_ext.agents.openai import OpenAIAgent +from openai import AsyncOpenAI +from pydantic import BaseModel + + +class FakeChunkDelta: + def __init__(self, content: Optional[str] = None, tool_calls: Optional[List[Any]] = None) -> None: + self.content = content + self.tool_calls = tool_calls + + +class FakeChunkChoice: + def __init__(self, delta: Optional[FakeChunkDelta] = None, finish_reason: Optional[str] = None) -> None: + self.delta = delta + self.finish_reason = finish_reason + self.index = 0 + + +class FakeChunk: + def __init__(self, id: str = "chunk-1", choices: Optional[List[FakeChunkChoice]] = None) -> None: + self.id = id + self.choices = choices or [] + + +class FakeToolCallFunction: + def __init__(self, name: str = "", arguments: str = "") -> None: + self.name = name + self.arguments = arguments + + +class FakeToolCall: + def __init__(self, id: str = "call-1", function: Optional[FakeToolCallFunction] = None) -> None: + self.id = id + self.type = "function" + self.function = function or FakeToolCallFunction() + + +def create_mock_openai_client() -> AsyncOpenAI: + """Create a mock OpenAI client for the Responses API.""" + client = AsyncMock(spec=AsyncOpenAI) + + async def mock_responses_create(**kwargs: Any) -> Any: + class MockResponse: + def __init__(self, output_text: str, id: str) -> None: + self.output_text = output_text + self.id = id + + if "tools" in kwargs and kwargs["tools"]: + return MockResponse(output_text='{"temperature": 72.5, "conditions": "sunny"}', id="resp-123") + return MockResponse(output_text="Hello world!", id="resp-abc") + + responses = MagicMock() + responses.create = AsyncMock(side_effect=mock_responses_create) + client.responses = responses + return client + + +@pytest.fixture +def mock_openai_client() -> AsyncOpenAI: + return create_mock_openai_client() + + +@pytest.fixture +def mock_error_client() -> AsyncOpenAI: + client = AsyncMock(spec=AsyncOpenAI) + beta = MagicMock() + client.beta = beta + beta.chat = MagicMock() + beta.chat.completions = MagicMock() + + async def mock_create_error(**kwargs: Any) -> None: + raise Exception("API Error") + + responses = MagicMock() + responses.create = AsyncMock(side_effect=mock_create_error) + client.responses = responses + return client + + +@pytest.fixture +def cancellation_token() -> CancellationToken: + return CancellationToken() + + +class WeatherResponse(BaseModel): + temperature: float + conditions: str + + +class GetWeatherArgs(BaseModel): + location: str + + +class WeatherTool(Tool): + def __init__(self) -> None: + self._name = "get_weather" + self._description = "Get the current weather in a location" + self._input_schema = GetWeatherArgs + self._output_schema = WeatherResponse + self._schema = ToolSchema( + name=self._name, + description=self._description, + parameters={ + "type": "object", + "properties": {"location": {"type": "string", "description": "The location to get weather for"}}, + "required": ["location"], + }, + ) + + @property + def name(self) -> str: + return self._name + + @property + def description(self) -> str: + return "Get the current weather in a location" + + @property + def schema(self) -> ToolSchema: + return self._schema + + def args_type(self) -> Type[BaseModel]: + return GetWeatherArgs + + def return_type(self) -> Type[Any]: + return WeatherResponse + + def state_type(self) -> Type[BaseModel] | None: + return None + + def return_value_as_string(self, value: Any) -> str: + if isinstance(value, dict): + return json.dumps(value) + return str(value) + + async def run_json( + self, args: Mapping[str, Any], cancellation_token: CancellationToken, call_id: str | None = None + ) -> Dict[str, Any]: + _ = GetWeatherArgs(**args) + return WeatherResponse(temperature=72.5, conditions="sunny").model_dump() + + async def load_state_json(self, state: Mapping[str, Any]) -> None: + pass + + async def save_state_json(self) -> Mapping[str, Any]: + return {} + + +@pytest.fixture +def weather_tool() -> WeatherTool: + return WeatherTool() + + +@pytest.fixture +def failing_tool() -> Tool: + tool = MagicMock(spec=Tool) + tool.name = "failing_tool" + tool.run_json = AsyncMock(side_effect=Exception("Tool execution failed")) + return tool + + +@pytest.fixture +def agent(mock_openai_client: AsyncOpenAI, weather_tool: WeatherTool) -> OpenAIAgent: + return OpenAIAgent( + name="assistant", + description="Test assistant using the Response API", + client=mock_openai_client, + model="gpt-4o", + instructions="You are a helpful AI assistant.", + tools=[weather_tool], + temperature=0.7, + max_output_tokens=1000, + store=True, + truncation="auto", + ) + + +@pytest.fixture +def json_mode_agent(mock_openai_client: AsyncOpenAI) -> OpenAIAgent: + return OpenAIAgent( + name="json_assistant", + description="JSON assistant", + client=mock_openai_client, + model="gpt-4o", + instructions="Return JSON responses", + json_mode=True, + ) + + +@pytest.fixture +def error_agent(mock_error_client: AsyncOpenAI) -> OpenAIAgent: + return OpenAIAgent( + name="error_assistant", + description="Assistant that generates errors", + client=mock_error_client, + model="gpt-4o", + instructions="You are a helpful AI assistant.", + ) + + +@pytest.mark.asyncio +async def test_basic_response(agent: OpenAIAgent, cancellation_token: CancellationToken) -> None: + """Test that the agent returns a basic text response from the Responses API.""" + message = TextMessage(source="user", content="Hello, how are you?") + response = await agent.on_messages([message], cancellation_token) + + assert response.chat_message is not None + assert isinstance(response.chat_message, TextMessage) + assert response.chat_message.content in ("Hello world!", '{"temperature": 72.5, "conditions": "sunny"}') + assert response.chat_message.source == "assistant" + + +@pytest.mark.asyncio +async def test_tool_calling(agent: OpenAIAgent, cancellation_token: CancellationToken) -> None: + """Test that the agent can call a tool and return the result using the Responses API.""" + + async def mock_run_json(self: Any, args: Dict[str, Any], cancellation_token: CancellationToken) -> Dict[str, Any]: + return {"temperature": 75.0, "conditions": "sunny and clear"} + + with patch.object(WeatherTool, "run_json", mock_run_json): + message = TextMessage(source="user", content="What's the weather in New York?") + + all_messages: List[Any] = [] + async for msg in agent.on_messages_stream([message], cancellation_token): + all_messages.append(msg) + + final_response = next((msg for msg in all_messages if hasattr(msg, "chat_message")), None) + assert final_response is not None + assert hasattr(final_response, "chat_message") + response_msg = cast(Response, final_response) + assert isinstance(response_msg.chat_message, TextMessage) + assert response_msg.chat_message.content in ( + '{"temperature": 75.0, "conditions": "sunny and clear"}', + '{"temperature": 72.5, "conditions": "sunny"}', + ) + + +@pytest.mark.asyncio +async def test_error_handling(error_agent: OpenAIAgent, cancellation_token: CancellationToken) -> None: + """Test that the agent returns an error message if the Responses API fails.""" + message = TextMessage(source="user", content="This will cause an error") + + all_messages: List[Any] = [] + async for msg in error_agent.on_messages_stream([message], cancellation_token): + all_messages.append(msg) + + final_response = next((msg for msg in all_messages if hasattr(msg, "chat_message")), None) + assert final_response is not None + assert isinstance(final_response.chat_message, TextMessage) + assert "Error generating response:" in final_response.chat_message.content + + +@pytest.mark.asyncio +async def test_state_management(agent: OpenAIAgent, cancellation_token: CancellationToken) -> None: + agent._last_response_id = "resp-123" # type: ignore + agent._message_history = [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there"}] # type: ignore + + state = await agent.save_state() + + new_agent = OpenAIAgent( + name="assistant2", + description="Test assistant 2", + client=agent._client, # type: ignore + model="gpt-4o", + instructions="You are a helpful AI assistant.", + ) + + await new_agent.load_state(state) + + assert new_agent._last_response_id == "resp-123" # type: ignore + assert len(new_agent._message_history) == 2 # type: ignore + assert new_agent._message_history[0]["role"] == "user" # type: ignore + assert new_agent._message_history[0]["content"] == "Hello" # type: ignore + + await new_agent.on_reset(cancellation_token) + assert new_agent._last_response_id is None # type: ignore + assert len(new_agent._message_history) == 0 # type: ignore + + +@pytest.mark.asyncio +async def test_convert_message_functions(agent: OpenAIAgent) -> None: + from autogen_ext.agents.openai._openai_agent import _convert_message_to_openai_message # type: ignore + + user_msg = TextMessage(content="Hello", source="user") + openai_user_msg = _convert_message_to_openai_message(user_msg) # type: ignore + assert openai_user_msg["role"] == "user" + assert openai_user_msg["content"] == "Hello" + + sys_msg = TextMessage(content="System prompt", source="system") + openai_sys_msg = _convert_message_to_openai_message(sys_msg) # type: ignore + assert openai_sys_msg["role"] == "system" + assert openai_sys_msg["content"] == "System prompt" + + assistant_msg = TextMessage(content="Assistant reply", source="assistant") + openai_assistant_msg = _convert_message_to_openai_message(assistant_msg) # type: ignore + assert openai_assistant_msg["role"] == "assistant" + assert openai_assistant_msg["content"] == "Assistant reply" + + text_msg = TextMessage(content="Plain text", source="other") + openai_text_msg = _convert_message_to_openai_message(text_msg) # type: ignore + assert openai_text_msg["role"] == "user" + assert openai_text_msg["content"] == "Plain text" + + +@pytest.mark.asyncio +async def test_tool_schema_conversion(agent: OpenAIAgent) -> None: + from autogen_ext.agents.openai._openai_agent import _convert_tool_to_function_schema # type: ignore + + tool_schema = _convert_tool_to_function_schema(agent._tool_map["get_weather"]) # type: ignore + + assert tool_schema["name"] == "get_weather" + assert "description" in tool_schema + assert "parameters" in tool_schema + assert tool_schema["parameters"]["type"] == "object" + assert "properties" in tool_schema["parameters"] + + +@pytest.mark.asyncio +async def test_on_messages_inner_messages(agent: OpenAIAgent, cancellation_token: CancellationToken) -> None: + class DummyMsg(BaseChatMessage): + type: str = "DummyMsg" + content: str = "dummy content" + + def __init__(self) -> None: + super().__init__(source="dummy") + + def to_model_message(self) -> UserMessage: + return UserMessage(content=self.content, source=self.source) + + def to_model_text(self) -> str: + return self.content + + def to_text(self) -> str: + return self.content + + dummy_inner = DummyMsg() + dummy_response = Response(chat_message=TextMessage(source="agent", content="hi"), inner_messages=None) + + async def fake_stream(*args: Any, **kwargs: Any) -> AsyncGenerator[Union[BaseChatMessage, Response], None]: + yield dummy_inner + yield dummy_response + + with patch.object(agent, "on_messages_stream", fake_stream): + response = await agent.on_messages([TextMessage(source="user", content="test")], cancellation_token) + assert response.chat_message is not None + assert isinstance(response.chat_message, TextMessage) + assert response.chat_message.content == "hi" + assert response.inner_messages is not None + assert dummy_inner in response.inner_messages + + +@pytest.mark.asyncio +async def test_build_api_params(agent: OpenAIAgent) -> None: + agent._last_response_id = None # type: ignore + params = agent._build_api_parameters([{"role": "user", "content": "hi"}]) # type: ignore + assert "previous_response_id" not in params + agent._last_response_id = "resp-456" # type: ignore + params = agent._build_api_parameters([{"role": "user", "content": "hi"}]) # type: ignore + assert params.get("previous_response_id") == "resp-456" + + assert "max_tokens" not in params + assert params.get("max_output_tokens") == 1000 + + assert params.get("store") is True + assert params.get("truncation") == "auto" + + agent._json_mode = True # type: ignore + params = agent._build_api_parameters([{"role": "user", "content": "hi"}]) # type: ignore + assert "text.format" not in params + assert params.get("text") == {"type": "json_object"} + + +@pytest.mark.asyncio +async def test_on_messages_previous_response_id(agent: OpenAIAgent, cancellation_token: CancellationToken) -> None: + message = TextMessage(source="user", content="hi") + response = await agent.on_messages([message], cancellation_token) + assert response.chat_message is not None + assert isinstance(response.chat_message, TextMessage) + message = TextMessage(source="user", content="hi") + response = await agent.on_messages([message], cancellation_token) + assert response.chat_message is not None + assert isinstance(response.chat_message, TextMessage) + + +@pytest.mark.asyncio +async def test_on_messages_stream(agent: OpenAIAgent, cancellation_token: CancellationToken) -> None: + dummy_response = Response(chat_message=TextMessage(source="agent", content="hi"), inner_messages=None) + + async def fake_stream(*args: Any, **kwargs: Any) -> AsyncGenerator[Response, None]: + yield dummy_response + + with patch.object(agent, "on_messages_stream", fake_stream): + resp = await agent.on_messages([TextMessage(source="user", content="hi")], cancellation_token) + assert isinstance(resp.chat_message, TextMessage) + assert resp.chat_message.content == "hi" + + +@pytest.mark.asyncio +async def test_component_serialization(agent: OpenAIAgent) -> None: + config = agent.dump_component() + config_dict: Any = None + if isinstance(config, dict): + config_dict = config + elif hasattr(config, "model_dump_json"): + config_dict = json.loads(config.model_dump_json()) + elif hasattr(config, "model_dump"): + config_dict = config.model_dump() + elif isinstance(config, str): + config_dict = json.loads(config) + else: + config_dict = {"name": agent.name, "description": agent.description} + + if isinstance(config_dict, dict) and "config" in config_dict: + config_dict = config_dict["config"] + + assert config_dict["name"] == "assistant" + assert config_dict["description"] == "Test assistant using the Response API" + assert config_dict["model"] == "gpt-4o" + assert config_dict["instructions"] == "You are a helpful AI assistant." + assert config_dict["temperature"] == 0.7 + assert config_dict["max_output_tokens"] == 1000 + assert config_dict["store"] is True + assert config_dict["truncation"] == "auto" + + +@pytest.mark.asyncio +async def test_from_config(agent: OpenAIAgent) -> None: + config = agent.dump_component() + + config_dict: Dict[str, Any] = {} + if hasattr(config, "model_dump_json"): + config_dict = json.loads(config.model_dump_json()) + elif isinstance(config, str): + config_dict = json.loads(config) + elif isinstance(config, dict): + config_dict = config + + if "tools" in config_dict and config_dict["tools"] is not None: + serialized_tools: List[Dict[str, Any]] = [] + tools_any: Any = config_dict["tools"] + if isinstance(tools_any, list): + tools_list: List[Any] = cast(List[Any], tools_any) # type: ignore[redundant-cast] + tools_count: int = len(tools_list) + for i in range(tools_count): + tool_any: Any = tools_list[i] + tool_dict: Dict[str, Any] = {} + if isinstance(tool_any, dict): + tool_dict = tool_any + elif tool_any is not None and isinstance(tool_any, object) and hasattr(tool_any, "model_dump"): + model_dump_any: Any = getattr(tool_any, "model_dump", None) + if callable(model_dump_any): + try: + result_any: Any = model_dump_any() + if isinstance(result_any, dict): + tool_dict = result_any + else: + tool_dict = {"provider": "unknown", "config": {}} + except Exception: + tool_dict = {"provider": "unknown", "config": {}} + else: + tool_dict = {"provider": "unknown", "config": {}} + else: + tool_dict = {"provider": "unknown", "config": {}} + serialized_tools.append(tool_dict) + config_dict["tools"] = serialized_tools + + with patch("openai.AsyncOpenAI"): + loaded_agent = OpenAIAgent.load_component(config_dict) + + assert loaded_agent.name == "assistant" + assert loaded_agent.description == "Test assistant using the Response API" + assert loaded_agent._model == "gpt-4o" # type: ignore + assert loaded_agent._instructions == "You are a helpful AI assistant." # type: ignore + assert loaded_agent._temperature == 0.7 # type: ignore + assert loaded_agent._max_output_tokens == 1000 # type: ignore + assert loaded_agent._store is True # type: ignore + assert loaded_agent._truncation == "auto" # type: ignore + + +@pytest.mark.asyncio +async def test_multimodal_message_response(agent: OpenAIAgent, cancellation_token: CancellationToken) -> None: + # Test that the multimodal message is converted to the correct format + img = Image.from_base64( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC" + ) + multimodal_message = MultiModalMessage(content=["Can you describe the content of this image?", img], source="user") + + # Patch client.responses.create to simulate image-capable output + async def mock_responses_create(**kwargs: Any) -> Any: + class MockResponse: + def __init__(self) -> None: + self.output_text = "I see a cat in the image." + self.id = "resp-image-001" + + return MockResponse() + + agent._client.responses.create = AsyncMock(side_effect=mock_responses_create) # type: ignore + + response = await agent.on_messages([multimodal_message], cancellation_token) + + assert response.chat_message is not None + assert isinstance(response.chat_message, TextMessage) + assert "cat" in response.chat_message.content.lower() diff --git a/python/packages/autogen-ext/tests/test_openai_assistant_agent.py b/python/packages/autogen-ext/tests/test_openai_assistant_agent.py index 43bd3447737a..4f75dc43a3f0 100644 --- a/python/packages/autogen-ext/tests/test_openai_assistant_agent.py +++ b/python/packages/autogen-ext/tests/test_openai_assistant_agent.py @@ -8,7 +8,7 @@ import aiofiles import pytest -from autogen_agentchat.messages import ChatMessage, TextMessage +from autogen_agentchat.messages import BaseChatMessage, TextMessage, ToolCallRequestEvent from autogen_core import CancellationToken from autogen_core.tools._base import BaseTool, Tool from autogen_ext.agents.openai import OpenAIAssistantAgent @@ -81,7 +81,7 @@ def __init__(self, id: str, text: str): class FakeCursorPage: - def __init__(self, data: List[ChatMessage | FakeMessage]) -> None: + def __init__(self, data: List[BaseChatMessage | FakeMessage]) -> None: self.data = data def has_next_page(self) -> bool: @@ -202,7 +202,7 @@ def agent(client: AsyncOpenAI) -> OpenAIAssistantAgent: return OpenAIAssistantAgent( name="assistant", instructions="Help the user with their task.", - model="gpt-4o-mini", + model="gpt-4.1-nano", description="OpenAI Assistant Agent", client=client, tools=tools, @@ -250,8 +250,7 @@ async def fake_async_aiofiles_open(*args: Any, **kwargs: Dict[str, Any]) -> Asyn message = TextMessage(source="user", content="What is the first sentence of the jungle scout book?") response = await agent.on_messages([message], cancellation_token) - assert response.chat_message.content is not None - assert isinstance(response.chat_message.content, str) + assert isinstance(response.chat_message, TextMessage) assert len(response.chat_message.content) > 0 await agent.delete_uploaded_files(cancellation_token) @@ -271,8 +270,7 @@ async def test_code_interpreter( message = TextMessage(source="user", content="I need to solve the equation `3x + 11 = 14`. Can you help me?") response = await agent.on_messages([message], cancellation_token) - assert response.chat_message.content is not None - assert isinstance(response.chat_message.content, str) + assert isinstance(response.chat_message, TextMessage) assert len(response.chat_message.content) > 0 assert "x = 1" in response.chat_message.content.lower() @@ -326,12 +324,11 @@ async def test_quiz_creation( response = await agent.on_messages([message], cancellation_token) # Check that the final response has non-empty inner messages (i.e. tool call events). - assert response.chat_message.content is not None - assert isinstance(response.chat_message.content, str) + assert isinstance(response.chat_message, TextMessage) assert len(response.chat_message.content) > 0 assert isinstance(response.inner_messages, list) # Ensure that at least one inner message has non-empty content. - assert any(hasattr(tool_msg, "content") and tool_msg.content for tool_msg in response.inner_messages) + assert any(isinstance(msg, ToolCallRequestEvent) for msg in response.inner_messages) await agent.delete_assistant(cancellation_token) @@ -339,8 +336,8 @@ async def test_quiz_creation( @pytest.mark.asyncio async def test_on_reset_behavior(client: AsyncOpenAI, cancellation_token: CancellationToken) -> None: # Arrange: Use the default behavior for reset. - thread = await client.beta.threads.create() - await client.beta.threads.messages.create( + thread = await client.beta.threads.create() # type: ignore[reportDeprecated] + await client.beta.threads.messages.create( # type: ignore[reportDeprecated] thread_id=thread.id, content="Hi, my name is John and I'm a software engineer. Use this information to help me.", role="user", @@ -349,7 +346,7 @@ async def test_on_reset_behavior(client: AsyncOpenAI, cancellation_token: Cancel agent = OpenAIAssistantAgent( name="assistant", instructions="Help the user with their task.", - model="gpt-4o-mini", + model="gpt-4.1-nano", description="OpenAI Assistant Agent", client=client, thread_id=thread.id, @@ -357,14 +354,14 @@ async def test_on_reset_behavior(client: AsyncOpenAI, cancellation_token: Cancel message1 = TextMessage(source="user", content="What is my name?") response1 = await agent.on_messages([message1], cancellation_token) - assert isinstance(response1.chat_message.content, str) + assert isinstance(response1.chat_message, TextMessage) assert "john" in response1.chat_message.content.lower() await agent.on_reset(cancellation_token) message2 = TextMessage(source="user", content="What is my name?") response2 = await agent.on_messages([message2], cancellation_token) - assert isinstance(response2.chat_message.content, str) + assert isinstance(response2.chat_message, TextMessage) assert "john" in response2.chat_message.content.lower() await agent.delete_assistant(cancellation_token) diff --git a/python/packages/autogen-ext/tests/test_websurfer_agent.py b/python/packages/autogen-ext/tests/test_websurfer_agent.py index 37423bfe6a50..371a8833be58 100644 --- a/python/packages/autogen-ext/tests/test_websurfer_agent.py +++ b/python/packages/autogen-ext/tests/test_websurfer_agent.py @@ -33,7 +33,7 @@ def emit(self, record: logging.LogRecord) -> None: record.msg = json.dumps( { "timestamp": ts, - "message": record.msg.model_dump(), + "message": record.msg.model_dump_json(indent=2), "type": record.msg.__class__.__name__, }, ) @@ -61,7 +61,7 @@ async def mock_create( @pytest.mark.asyncio async def test_run_websurfer(monkeypatch: pytest.MonkeyPatch) -> None: - model = "gpt-4o-2024-05-13" + model = "gpt-4.1-nano-2025-04-14" chat_completions = [ ChatCompletion( id="id2", @@ -149,7 +149,7 @@ async def test_run_websurfer(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.mark.asyncio async def test_run_websurfer_declarative(monkeypatch: pytest.MonkeyPatch) -> None: - model = "gpt-4o-2024-05-13" + model = "gpt-4.1-nano-2025-04-14" chat_completions = [ ChatCompletion( id="id1", diff --git a/python/packages/autogen-ext/tests/test_worker_runtime.py b/python/packages/autogen-ext/tests/test_worker_runtime.py index dede306853ed..ec57f187e821 100644 --- a/python/packages/autogen-ext/tests/test_worker_runtime.py +++ b/python/packages/autogen-ext/tests/test_worker_runtime.py @@ -577,6 +577,139 @@ async def test_grpc_max_message_size() -> None: await host.stop() +@pytest.mark.grpc +@pytest.mark.asyncio +async def test_agent_type_register_instance() -> None: + host_address = "localhost:50051" + agent1_id = AgentId(type="name", key="default") + agentdup_id = AgentId(type="name", key="default") + agent2_id = AgentId(type="name", key="notdefault") + host = GrpcWorkerAgentRuntimeHost(address=host_address) + host.start() + + worker = GrpcWorkerAgentRuntime(host_address=host_address) + agent1 = NoopAgent() + agent2 = NoopAgent() + agentdup = NoopAgent() + await worker.start() + + await worker.register_agent_instance(agent1, agent_id=agent1_id) + await worker.register_agent_instance(agent2, agent_id=agent2_id) + + with pytest.raises(ValueError): + await worker.register_agent_instance(agentdup, agent_id=agentdup_id) + + assert await worker.try_get_underlying_agent_instance(agent1_id, type=NoopAgent) == agent1 + assert await worker.try_get_underlying_agent_instance(agent2_id, type=NoopAgent) == agent2 + + await worker.stop() + await host.stop() + + +@pytest.mark.grpc +@pytest.mark.asyncio +async def test_agent_type_register_instance_different_types() -> None: + host_address = "localhost:50051" + agent1_id = AgentId(type="name", key="noop") + agent2_id = AgentId(type="name", key="loopback") + host = GrpcWorkerAgentRuntimeHost(address=host_address) + host.start() + + worker = GrpcWorkerAgentRuntime(host_address=host_address) + agent1 = NoopAgent() + agent2 = LoopbackAgent() + await worker.start() + + await worker.register_agent_instance(agent1, agent_id=agent1_id) + with pytest.raises(ValueError): + await worker.register_agent_instance(agent2, agent_id=agent2_id) + + await worker.stop() + await host.stop() + + +@pytest.mark.grpc +@pytest.mark.asyncio +async def test_register_instance_factory() -> None: + host_address = "localhost:50051" + agent1_id = AgentId(type="name", key="default") + host = GrpcWorkerAgentRuntimeHost(address=host_address) + host.start() + + worker = GrpcWorkerAgentRuntime(host_address=host_address) + agent1 = NoopAgent() + await worker.start() + + await agent1.register_instance(runtime=worker, agent_id=agent1_id) + + with pytest.raises(ValueError): + await NoopAgent.register(runtime=worker, type="name", factory=lambda: NoopAgent()) + + await worker.stop() + await host.stop() + + +@pytest.mark.grpc +@pytest.mark.asyncio +async def test_instance_factory_messaging() -> None: + host_address = "localhost:50051" + loopback_agent_id = AgentId(type="dm_agent", key="dm_agent") + cascading_agent_id = AgentId(type="instance_agent", key="instance_agent") + host = GrpcWorkerAgentRuntimeHost(address=host_address) + host.start() + + worker = GrpcWorkerAgentRuntime(host_address=host_address) + cascading_agent = CascadingAgent(max_rounds=5) + loopback_agent = LoopbackAgent() + await worker.start() + + await loopback_agent.register_instance(worker, agent_id=loopback_agent_id) + resp = await worker.send_message(message=ContentMessage(content="Hello!"), recipient=loopback_agent_id) + assert resp == ContentMessage(content="Hello!") + + await cascading_agent.register_instance(worker, agent_id=cascading_agent_id) + await CascadingAgent.register(worker, "factory_agent", lambda: CascadingAgent(max_rounds=5)) + + # instance_agent will publish a message that factory_agent will pick up + for i in range(5): + await worker.publish_message( + CascadingMessageType(round=i + 1), TopicId(type="instance_agent", source="instance_agent") + ) + await asyncio.sleep(2) + + agent = await worker.try_get_underlying_agent_instance(AgentId("factory_agent", "default"), CascadingAgent) + assert agent.num_calls == 4 + assert cascading_agent.num_calls == 5 + + await worker.stop() + await host.stop() + + +# GrpcWorkerAgentRuntimeHost eats exceptions in the main loop +# @pytest.mark.grpc +# @pytest.mark.asyncio +# async def test_agent_type_register_instance_publish_new_source() -> None: +# host_address = "localhost:50056" +# agent_id = AgentId(type="name", key="default") +# agent1 = LoopbackAgent() +# host = GrpcWorkerAgentRuntimeHost(address=host_address) +# host.start() +# worker = GrpcWorkerAgentRuntime(host_address=host_address) +# await worker.start() +# publisher = GrpcWorkerAgentRuntime(host_address=host_address) +# publisher.add_message_serializer(try_get_known_serializers_for_type(MessageType)) +# await publisher.start() + +# await agent1.register_instance(worker, agent_id=agent_id) +# await worker.add_subscription(TypeSubscription("notdefault", "name")) + +# with pytest.raises(RuntimeError): +# await worker.publish_message(MessageType(), TopicId("notdefault", "notdefault")) +# await asyncio.sleep(2) + +# await worker.stop() +# await host.stop() + if __name__ == "__main__": os.environ["GRPC_VERBOSITY"] = "DEBUG" os.environ["GRPC_TRACE"] = "all" diff --git a/python/packages/autogen-ext/tests/tools/azure/conftest.py b/python/packages/autogen-ext/tests/tools/azure/conftest.py new file mode 100644 index 000000000000..6d3a8569f127 --- /dev/null +++ b/python/packages/autogen-ext/tests/tools/azure/conftest.py @@ -0,0 +1,223 @@ +"""Test fixtures for Azure AI Search tool tests.""" + +import warnings +from typing import Any, Dict, Iterator, List, Protocol, TypeVar, Union +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from autogen_core import ComponentModel + +T = TypeVar("T") + +try: + from azure.core.credentials import AzureKeyCredential, TokenCredential + + azure_sdk_available = True +except ImportError: + azure_sdk_available = False + +skip_if_no_azure_sdk = pytest.mark.skipif( + not azure_sdk_available, reason="Azure SDK components (azure-search-documents, azure-identity) not available" +) + + +class AccessTokenProtocol(Protocol): + """Protocol matching Azure AccessToken.""" + + token: str + expires_on: int + + +class MockAccessToken: + """Mock implementation of AccessToken.""" + + def __init__(self, token: str, expires_on: int) -> None: + self.token = token + self.expires_on = expires_on + + +class MockAzureKeyCredential: + """Mock implementation of AzureKeyCredential.""" + + def __init__(self, key: str) -> None: + self.key = key + + +class MockTokenCredential: + """Mock implementation of TokenCredential for testing.""" + + def get_token( + self, + *scopes: str, + claims: str | None = None, + tenant_id: str | None = None, + enable_cae: bool = False, + **kwargs: Any, + ) -> AccessTokenProtocol: + """Mock get_token method that implements TokenCredential protocol.""" + return MockAccessToken("mock-token", 12345) + + +CredentialType = Union[ + AzureKeyCredential, # pyright: ignore [reportPossiblyUnboundVariable] + TokenCredential, # pyright: ignore [reportPossiblyUnboundVariable] + MockAzureKeyCredential, + MockTokenCredential, + Any, +] + +needs_azure_sdk = pytest.mark.skipif(not azure_sdk_available, reason="Azure SDK not available") + +warnings.filterwarnings( + "ignore", + message="Type google.*uses PyType_Spec with a metaclass that has custom tp_new", + category=DeprecationWarning, +) + + +@pytest.fixture +def mock_vectorized_query() -> MagicMock: + """Create a mock VectorizedQuery for testing.""" + if azure_sdk_available: + from azure.search.documents.models import VectorizedQuery + + return MagicMock(spec=VectorizedQuery) + else: + return MagicMock() + + +@pytest.fixture +def test_config() -> ComponentModel: + """Return a test configuration for the Azure AI Search tool.""" + return ComponentModel( + provider="autogen_ext.tools.azure.MockAzureAISearchTool", + config={ + "name": "TestAzureSearch", + "description": "Test Azure AI Search Tool", + "endpoint": "https://test-search-service.search.windows.net", + "index_name": "test-index", + "api_version": "2023-10-01-Preview", + "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"}, # pyright: ignore [reportPossiblyUnboundVariable] + "query_type": "keyword", + "search_fields": ["content", "title"], + "select_fields": ["id", "content", "title", "source"], + "top": 5, + }, + ) + + +@pytest.fixture +def keyword_config() -> ComponentModel: + """Return a keyword search configuration.""" + return ComponentModel( + provider="autogen_ext.tools.azure.MockAzureAISearchTool", + config={ + "name": "KeywordSearch", + "description": "Keyword search tool", + "endpoint": "https://test-search-service.search.windows.net", + "index_name": "test-index", + "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"}, # pyright: ignore [reportPossiblyUnboundVariable] + "query_type": "keyword", + "search_fields": ["content", "title"], + "select_fields": ["id", "content", "title", "source"], + }, + ) + + +@pytest.fixture +def vector_config() -> ComponentModel: + """Create a test configuration for vector search.""" + return ComponentModel( + provider="autogen_ext.tools.azure.MockAzureAISearchTool", + config={ + "name": "VectorSearch", + "description": "Vector search tool", + "endpoint": "https://test-search-service.search.windows.net", + "index_name": "test-index", + "api_version": "2023-10-01-Preview", + "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"}, # pyright: ignore [reportPossiblyUnboundVariable] + "query_type": "vector", + "vector_fields": ["embedding"], + "select_fields": ["id", "content", "title", "source"], + "top": 5, + }, + ) + + +@pytest.fixture +def hybrid_config() -> ComponentModel: + """Create a test configuration for hybrid search.""" + return ComponentModel( + provider="autogen_ext.tools.azure.MockAzureAISearchTool", + config={ + "name": "HybridSearch", + "description": "Hybrid search tool", + "endpoint": "https://test-search-service.search.windows.net", + "index_name": "test-index", + "api_version": "2023-10-01-Preview", + "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"}, # pyright: ignore [reportPossiblyUnboundVariable] + "query_type": "keyword", + "search_fields": ["content", "title"], + "vector_fields": ["embedding"], + "select_fields": ["id", "content", "title", "source"], + "top": 5, + }, + ) + + +@pytest.fixture +def mock_search_response() -> List[Dict[str, Any]]: + """Create a mock search response.""" + return [ + { + "@search.score": 0.95, + "id": "doc1", + "content": "This is the first document content", + "title": "Document 1", + "source": "test-source-1", + }, + { + "@search.score": 0.85, + "id": "doc2", + "content": "This is the second document content", + "title": "Document 2", + "source": "test-source-2", + }, + ] + + +class AsyncIterator: + """Async iterator for testing.""" + + def __init__(self, items: List[Dict[str, Any]]) -> None: + self.items = items.copy() + + def __aiter__(self) -> "AsyncIterator": + return self + + async def __anext__(self) -> Dict[str, Any]: + if not self.items: + raise StopAsyncIteration + return self.items.pop(0) + + async def get_count(self) -> int: + """Return count of items.""" + return len(self.items) + + +@pytest.fixture +def mock_search_client(mock_search_response: List[Dict[str, Any]]) -> Iterator[MagicMock]: + """Create a mock search client for testing, with the patch active.""" + mock_client_instance = MagicMock() + mock_client_instance.__aenter__ = AsyncMock(return_value=mock_client_instance) + mock_client_instance.__aexit__ = AsyncMock(return_value=None) + + search_results_iterator = AsyncIterator(mock_search_response) + mock_client_instance.search = MagicMock(return_value=search_results_iterator) + + patch_target = "autogen_ext.tools.azure._ai_search.SearchClient" + patcher = patch(patch_target, return_value=mock_client_instance) + + patcher.start() + yield mock_client_instance + patcher.stop() diff --git a/python/packages/autogen-ext/tests/tools/azure/test_ai_search_config.py b/python/packages/autogen-ext/tests/tools/azure/test_ai_search_config.py new file mode 100644 index 000000000000..ddcc26e36869 --- /dev/null +++ b/python/packages/autogen-ext/tests/tools/azure/test_ai_search_config.py @@ -0,0 +1,297 @@ +from typing import Any, Dict, cast + +import pytest +from autogen_ext.tools.azure._config import AzureAISearchConfig, QueryTypeLiteral +from azure.core.credentials import AzureKeyCredential +from pydantic import ValidationError + +from tests.tools.azure.conftest import azure_sdk_available + +skip_if_no_azure_sdk = pytest.mark.skipif( + not azure_sdk_available, reason="Azure SDK components (azure-search-documents, azure-identity) not available" +) + +# ===================================== +# Basic Configuration Tests +# ===================================== + + +def test_basic_config_creation() -> None: + """Test that a basic valid configuration can be created.""" + config = AzureAISearchConfig( + name="test_tool", + endpoint="https://test-search.search.windows.net", + index_name="test-index", + credential=AzureKeyCredential("test-key"), + ) + + assert config.name == "test_tool" + assert config.endpoint == "https://test-search.search.windows.net" + assert config.index_name == "test-index" + assert isinstance(config.credential, AzureKeyCredential) + assert config.query_type == "simple" # default value + + +def test_endpoint_validation() -> None: + """Test that endpoint validation works correctly.""" + valid_endpoints = ["https://test.search.windows.net", "http://localhost:8080"] + + for endpoint in valid_endpoints: + config = AzureAISearchConfig( + name="test_tool", + endpoint=endpoint, + index_name="test-index", + credential=AzureKeyCredential("test-key"), + ) + assert config.endpoint == endpoint + + invalid_endpoints = [ + "test.search.windows.net", + "ftp://test.search.windows.net", + "", + ] + + for endpoint in invalid_endpoints: + with pytest.raises(ValidationError) as exc: + AzureAISearchConfig( + name="test_tool", + endpoint=endpoint, + index_name="test-index", + credential=AzureKeyCredential("test-key"), + ) + assert "endpoint must be a valid URL" in str(exc.value) + + +def test_top_validation() -> None: + """Test validation of top parameter.""" + valid_tops = [1, 5, 10, 100] + + for top in valid_tops: + config = AzureAISearchConfig( + name="test_tool", + endpoint="https://test.search.windows.net", + index_name="test-index", + credential=AzureKeyCredential("test-key"), + top=top, + ) + assert config.top == top + + invalid_tops = [0, -1, -10] + + for top in invalid_tops: + with pytest.raises(ValidationError) as exc: + AzureAISearchConfig( + name="test_tool", + endpoint="https://test.search.windows.net", + index_name="test-index", + credential=AzureKeyCredential("test-key"), + top=top, + ) + assert "top must be a positive integer" in str(exc.value) + + +# ===================================== +# Query Type Tests +# ===================================== + + +def test_query_type_normalization() -> None: + """Test that query_type normalization works correctly.""" + standard_query_types = { + "simple": "simple", + "full": "full", + "semantic": "semantic", + "vector": "vector", + } + + for input_type, expected_type in standard_query_types.items(): + config_args: Dict[str, Any] = { + "name": "test_tool", + "endpoint": "https://test.search.windows.net", + "index_name": "test-index", + "credential": AzureKeyCredential("test-key"), + "query_type": cast(QueryTypeLiteral, input_type), + } + + if input_type == "semantic": + config_args["semantic_config_name"] = "my-semantic-config" + elif input_type == "vector": + config_args["vector_fields"] = ["content_vector"] + + config = AzureAISearchConfig(**config_args) + assert config.query_type == expected_type + + with pytest.raises(ValidationError) as exc: + AzureAISearchConfig( + name="test_tool", + endpoint="https://test.search.windows.net", + index_name="test-index", + credential=AzureKeyCredential("test-key"), + query_type=cast(Any, "invalid_type"), + ) + assert "Input should be" in str(exc.value) + + +def test_semantic_config_validation() -> None: + """Test validation of semantic configuration.""" + config = AzureAISearchConfig( + name="test_tool", + endpoint="https://test.search.windows.net", + index_name="test-index", + credential=AzureKeyCredential("test-key"), + query_type=cast(QueryTypeLiteral, "semantic"), + semantic_config_name="my-semantic-config", + ) + assert config.query_type == "semantic" + assert config.semantic_config_name == "my-semantic-config" + + with pytest.raises(ValidationError) as exc: + AzureAISearchConfig( + name="test_tool", + endpoint="https://test.search.windows.net", + index_name="test-index", + credential=AzureKeyCredential("test-key"), + query_type=cast(QueryTypeLiteral, "semantic"), + ) + assert "semantic_config_name must be provided" in str(exc.value) + + +def test_vector_fields_validation() -> None: + """Test validation of vector fields for vector search.""" + config = AzureAISearchConfig( + name="test_tool", + endpoint="https://test.search.windows.net", + index_name="test-index", + credential=AzureKeyCredential("test-key"), + query_type=cast(QueryTypeLiteral, "vector"), + vector_fields=["content_vector"], + ) + assert config.query_type == "vector" + assert config.vector_fields == ["content_vector"] + + +# ===================================== +# Embedding Configuration Tests +# ===================================== + + +def test_azure_openai_endpoint_validation() -> None: + """Test validation of Azure OpenAI endpoint for client-side embeddings.""" + config = AzureAISearchConfig( + name="test_tool", + endpoint="https://test.search.windows.net", + index_name="test-index", + credential=AzureKeyCredential("test-key"), + embedding_provider="azure_openai", + embedding_model="text-embedding-ada-002", + openai_endpoint="https://test.openai.azure.com", + ) + assert config.embedding_provider == "azure_openai" + assert config.embedding_model == "text-embedding-ada-002" + assert config.openai_endpoint == "https://test.openai.azure.com" + + with pytest.raises(ValidationError) as exc: + AzureAISearchConfig( + name="test_tool", + endpoint="https://test.search.windows.net", + index_name="test-index", + credential=AzureKeyCredential("test-key"), + embedding_provider="azure_openai", + embedding_model="text-embedding-ada-002", + ) + assert "openai_endpoint must be provided for azure_openai" in str(exc.value) + + config = AzureAISearchConfig( + name="test_tool", + endpoint="https://test.search.windows.net", + index_name="test-index", + credential=AzureKeyCredential("test-key"), + embedding_provider="openai", + embedding_model="text-embedding-ada-002", + ) + assert config.embedding_provider == "openai" + assert config.embedding_model == "text-embedding-ada-002" + assert config.openai_endpoint is None + + +# ===================================== +# Credential and Serialization Tests +# ===================================== + + +def test_credential_validation() -> None: + """Test credential validation scenarios.""" + config = AzureAISearchConfig( + name="test_tool", + endpoint="https://test.search.windows.net", + index_name="test-index", + credential=AzureKeyCredential("test-key"), + ) + assert isinstance(config.credential, AzureKeyCredential) + assert config.credential.key == "test-key" + + if azure_sdk_available: + from azure.core.credentials import AccessToken + from azure.core.credentials_async import AsyncTokenCredential + + class TestTokenCredential(AsyncTokenCredential): + async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken: + return AccessToken("test-token", 12345) + + async def close(self) -> None: + pass + + async def __aenter__(self) -> "TestTokenCredential": + return self + + async def __aexit__(self, *args: Any) -> None: + await self.close() + + config = AzureAISearchConfig( + name="test", + endpoint="https://endpoint", + index_name="index", + credential=TestTokenCredential(), + ) + assert isinstance(config.credential, AsyncTokenCredential) + + +def test_model_dump_scenarios() -> None: + """Test all model_dump scenarios to ensure full code coverage.""" + config = AzureAISearchConfig( + name="test", + endpoint="https://endpoint", + index_name="index", + credential=AzureKeyCredential("key"), + ) + result = config.model_dump() + assert isinstance(result["credential"], AzureKeyCredential) + assert result["credential"].key == "key" + + if azure_sdk_available: + from azure.core.credentials import AccessToken + from azure.core.credentials_async import AsyncTokenCredential + + class TestTokenCredential(AsyncTokenCredential): + async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken: + return AccessToken("test-token", 12345) + + async def close(self) -> None: + pass + + async def __aenter__(self) -> "TestTokenCredential": + return self + + async def __aexit__(self, *args: Any) -> None: + await self.close() + + config = AzureAISearchConfig( + name="test", + endpoint="https://endpoint", + index_name="index", + credential=TestTokenCredential(), + ) + result = config.model_dump() + assert isinstance(result["credential"], AsyncTokenCredential) + else: + pytest.skip("Skipping TokenCredential test - Azure SDK not available") diff --git a/python/packages/autogen-ext/tests/tools/azure/test_ai_search_tool.py b/python/packages/autogen-ext/tests/tools/azure/test_ai_search_tool.py new file mode 100644 index 000000000000..67b1d357fd50 --- /dev/null +++ b/python/packages/autogen-ext/tests/tools/azure/test_ai_search_tool.py @@ -0,0 +1,971 @@ +"""Tests for Azure AI Search tool.""" + +import asyncio +from collections.abc import Generator +from typing import Any, Dict, List +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from autogen_core import CancellationToken +from autogen_ext.tools.azure import ( + AzureAISearchConfig, + AzureAISearchTool, + SearchResult, + SearchResults, +) +from autogen_ext.tools.azure._ai_search import BaseAzureAISearchTool +from autogen_ext.tools.azure._config import DEFAULT_API_VERSION +from azure.core.credentials import AzureKeyCredential +from azure.core.credentials_async import AsyncTokenCredential +from azure.core.exceptions import HttpResponseError, ResourceNotFoundError +from pydantic import BaseModel, Field, ValidationError + +MOCK_ENDPOINT = "https://test-search.search.windows.net" +MOCK_INDEX = "test-index" +MOCK_API_KEY = "test-key" +MOCK_CREDENTIAL = AzureKeyCredential(MOCK_API_KEY) + + +class MockAsyncTokenCredential(AsyncTokenCredential): + """Mock async token credential for testing.""" + + async def get_token(self, *scopes: str, **kwargs: Any) -> Any: + return "mock-token" + + async def close(self) -> None: + pass + + async def __aexit__(self, exc_type: Any = None, exc_val: Any = None, exc_tb: Any = None) -> None: + await self.close() + + +@pytest.fixture +def search_config() -> AzureAISearchConfig: + """Fixture for basic search configuration.""" + return AzureAISearchConfig( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + description="Test search tool", + ) + + +@pytest.fixture +def mock_search_client() -> Generator[AsyncMock, None, None]: + """Fixture for mocked search client.""" + with patch("azure.search.documents.aio.SearchClient", autospec=True) as mock: + mock_client = AsyncMock() + mock.return_value = mock_client + yield mock_client + + +@pytest.fixture +def mock_search_results() -> List[Dict[str, Any]]: + """Fixture for mock search results.""" + return [ + { + "id": "1", + "content": "Test content", + "@search.score": 0.8, + } + ] + + +class TestSearchQuery(BaseModel): + """Test model for query validation.""" + + query: str = Field(min_length=1) + + +@pytest.mark.asyncio +async def test_search_query_model() -> None: + """Test SearchQuery model validation.""" + query = TestSearchQuery(query="test query") + assert query.query == "test query" + + with pytest.raises(ValidationError): + TestSearchQuery(query="") + + +@pytest.mark.asyncio +async def test_search_result_model() -> None: + """Test SearchResult model.""" + result = SearchResult(score=0.8, content={"title": "Test", "text": "Content"}, metadata={"@search.score": 0.8}) + assert result.score == 0.8 + assert result.content["title"] == "Test" + assert result.metadata["@search.score"] == 0.8 + + +@pytest.mark.asyncio +async def test_search_results_model() -> None: + """Test SearchResults model.""" + results = SearchResults( + results=[ + SearchResult(score=0.8, content={"title": "Test1"}, metadata={"@search.score": 0.8}), + SearchResult(score=0.6, content={"title": "Test2"}, metadata={"@search.score": 0.6}), + ] + ) + assert len(results.results) == 2 + assert results.results[0].score == 0.8 + assert results.results[1].content["title"] == "Test2" + + +@pytest.mark.asyncio +async def test_create_full_text_search() -> None: + """Test creation of full text search tool.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + search_fields=["content"], + query_type="simple", + ) + assert tool.name == "test-search" + assert tool.search_config.query_type == "simple" + assert tool.search_config.search_fields == ["content"] + + with pytest.raises(ValueError, match="semantic_config_name is required"): + AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + query_type="semantic", + ) + + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + query_type="semantic", + semantic_config_name="default", + ) + assert tool.search_config.query_type == "semantic" + assert tool.search_config.semantic_config_name == "default" + + +@pytest.mark.asyncio +async def test_create_vector_search() -> None: + """Test creation of vector search tool.""" + tool = AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + ) + assert tool.search_config.query_type == "vector" + assert tool.search_config.vector_fields == ["embedding"] + + with pytest.raises(ValueError, match="openai_endpoint is required"): + AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + embedding_provider="azure_openai", + embedding_model="text-embedding-ada-002", + ) + + +@pytest.mark.asyncio +async def test_create_hybrid_search() -> None: + """Test creation of hybrid search tool.""" + tool = AzureAISearchTool.create_hybrid_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + search_fields=["content"], + ) + assert tool.search_config.vector_fields == ["embedding"] + assert tool.search_config.search_fields == ["content"] + + tool = AzureAISearchTool.create_hybrid_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + search_fields=["content"], + query_type="semantic", + semantic_config_name="default", + ) + assert tool.search_config.query_type == "semantic" + assert tool.search_config.semantic_config_name == "default" + + +@pytest.mark.asyncio +async def test_search_execution(mock_search_client: AsyncMock, mock_search_results: List[Dict[str, Any]]) -> None: + """Test search execution with mocked client.""" + mock_search_client.search.return_value.__aiter__.return_value = mock_search_results + + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + with patch.object(tool, "_get_client", return_value=mock_search_client): + results = await tool.run("test query") + + assert len(results.results) == 1 + assert results.results[0].score == 0.8 + assert results.results[0].content["content"] == "Test content" + + mock_search_client.search.assert_called_once() + call_kwargs = mock_search_client.search.call_args[1] + assert call_kwargs["search_text"] == "test query" + + +@pytest.mark.asyncio +async def test_search_with_caching(mock_search_client: AsyncMock, mock_search_results: List[Dict[str, Any]]) -> None: + """Test search caching functionality.""" + mock_search_client.search.return_value.__aiter__.return_value = mock_search_results + + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + enable_caching=True, + cache_ttl_seconds=300, + ) + + with patch.object(tool, "_get_client", return_value=mock_search_client): + await tool.run("test query") + assert mock_search_client.search.call_count == 1 + + await tool.run("test query") + assert mock_search_client.search.call_count == 1 + + +@pytest.mark.asyncio +async def test_error_handling(mock_search_client: AsyncMock) -> None: + """Test error handling in search execution.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + with patch.object(tool, "_get_client", return_value=mock_search_client): + mock_search_client.search.side_effect = ResourceNotFoundError("Index not found") + with pytest.raises(ValueError, match="Index.*not found"): + await tool.run("test query") + + mock_search_client.search.side_effect = HttpResponseError(status_code=401, message="Unauthorized") + with pytest.raises(ValueError, match="Authentication failed"): + await tool.run("test query") + + mock_search_client.search.side_effect = HttpResponseError(status_code=500, message="Internal server error") + with pytest.raises(ValueError, match="Error from Azure AI Search"): + await tool.run("test query") + + +@pytest.mark.asyncio +async def test_embedding_provider_mixin() -> None: + """Test the embedding provider functionality.""" + with patch("openai.AsyncAzureOpenAI") as mock_azure_openai: + mock_client = AsyncMock() + mock_azure_openai.return_value = mock_client + mock_client.embeddings.create.return_value.data = [MagicMock(embedding=[0.1, 0.2, 0.3])] + + tool = AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + embedding_provider="azure_openai", + embedding_model="text-embedding-ada-002", + openai_endpoint="https://test.openai.azure.com", + openai_api_key="test-key", + ) + + embedding = await tool._get_embedding("test query") # pyright: ignore[reportPrivateUsage] + assert len(embedding) == 3 + assert embedding == [0.1, 0.2, 0.3] + + mock_client.embeddings.create.assert_called_once_with(model="text-embedding-ada-002", input="test query") + + +@pytest.mark.asyncio +async def test_credential_processing() -> None: + """Test credential processing logic.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + assert isinstance(tool.search_config.credential, AzureKeyCredential) + assert tool.search_config.credential.key == MOCK_API_KEY + + mock_async_credential = MockAsyncTokenCredential() + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=mock_async_credential, + ) + assert isinstance(tool.search_config.credential, AsyncTokenCredential) + + with pytest.raises(ValueError, match="Invalid configuration"): + AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential={"api_key": "test-key"}, # type: ignore + ) + + +@pytest.mark.asyncio +async def test_return_value_as_string() -> None: + """Test the string representation of search results.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + results = SearchResults( + results=[ + SearchResult(score=0.8, content={"title": "Test1", "text": "Content1"}, metadata={"@search.score": 0.8}), + SearchResult(score=0.6, content={"title": "Test2", "text": "Content2"}, metadata={"@search.score": 0.6}), + ] + ) + result_str = tool.return_value_as_string(results) + assert "Result 1 (Score: 0.80)" in result_str + assert "Result 2 (Score: 0.60)" in result_str + assert "Test1" in result_str + assert "Content2" in result_str + + empty_results = SearchResults(results=[]) + assert tool.return_value_as_string(empty_results) == "No results found." + + +@pytest.mark.asyncio +async def test_schema() -> None: + """Test tool schema generation.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + schema = tool.schema + assert schema["name"] == "test-search" + assert "description" in schema + assert "parameters" in schema + assert "required" in schema["parameters"] + assert schema["parameters"]["type"] == "object" + assert "query" in schema["parameters"]["properties"] + assert schema["parameters"]["required"] == ["query"] + + +@pytest.mark.asyncio +async def test_vector_search_execution( + mock_search_client: AsyncMock, mock_search_results: List[Dict[str, Any]] +) -> None: + """Test vector search execution.""" + mock_search_client.search.return_value.__aiter__.return_value = mock_search_results + + tool = AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + embedding_provider="azure_openai", + embedding_model="text-embedding-ada-002", + openai_endpoint="https://test.openai.azure.com", + openai_api_key="test-key", + ) + + mock_embedding = [0.1, 0.2, 0.3] + with ( + patch.object(tool, "_get_embedding", return_value=mock_embedding), + patch.object(tool, "_get_client", return_value=mock_search_client), + ): + results = await tool.run("test query") + assert len(results.results) == 1 + mock_search_client.search.assert_called_once() + + +@pytest.mark.asyncio +async def test_cancellation() -> None: + """Test search cancellation.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + token = CancellationToken() + token.cancel() + with pytest.raises(asyncio.CancelledError): + await tool.run("test query", cancellation_token=token) + + +@pytest.mark.asyncio +async def test_invalid_query_format() -> None: + """Test invalid query format handling.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + with pytest.raises(ValueError, match="Invalid search query format"): + await tool.run({"invalid": "format"}) + + +@pytest.mark.asyncio +async def test_client_cleanup() -> None: + """Test client cleanup.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + mock_client = AsyncMock() + tool._client = mock_client # pyright: ignore[reportPrivateUsage] + await tool.close() + + mock_client.close.assert_called_once() + assert tool._client is None # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_config_validation() -> None: + """Test configuration validation.""" + with pytest.raises(ValueError, match="vector_fields must contain at least one field"): + AzureAISearchTool._validate_config( # pyright: ignore[reportPrivateUsage] + { + "name": "test-search", + "endpoint": MOCK_ENDPOINT, + "index_name": MOCK_INDEX, + "credential": MOCK_CREDENTIAL, + }, + "vector", + ) + + with pytest.raises(ValueError, match="vector_fields must contain at least one field"): + AzureAISearchTool._validate_config( # pyright: ignore[reportPrivateUsage] + { + "name": "test-search", + "endpoint": MOCK_ENDPOINT, + "index_name": MOCK_INDEX, + "credential": MOCK_CREDENTIAL, + "search_fields": ["content"], + }, + "hybrid", + ) + + +@pytest.mark.asyncio +async def test_openai_embedding_provider() -> None: + """Test OpenAI embedding provider.""" + with patch("openai.AsyncOpenAI") as mock_openai: + mock_client = AsyncMock() + mock_openai.return_value = mock_client + mock_client.embeddings.create.return_value.data = [MagicMock(embedding=[0.1, 0.2, 0.3])] + + tool = AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + embedding_provider="openai", + embedding_model="text-embedding-ada-002", + openai_api_key="test-key", + ) + + embedding = await tool._get_embedding("test query") # pyright: ignore[reportPrivateUsage] + assert len(embedding) == 3 + assert embedding == [0.1, 0.2, 0.3] + + +@pytest.mark.asyncio +async def test_embedding_provider_error_handling() -> None: + """Test error handling in embedding providers.""" + with pytest.raises(ValueError, match="openai_endpoint is required when embedding_provider is 'azure_openai'"): + AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + embedding_provider="azure_openai", + embedding_model="text-embedding-ada-002", + openai_api_version="2023-11-01", + ) + + tool = AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + embedding_provider="azure_openai", + embedding_model="text-embedding-ada-002", + openai_endpoint="https://test.openai.azure.com", + openai_api_version="2023-11-01", + ) + with patch("azure.identity.DefaultAzureCredential") as mock_credential: + mock_credential.return_value.get_token.return_value = None + with pytest.raises(ValueError, match="Failed to acquire token using DefaultAzureCredential for Azure OpenAI"): + await tool._get_embedding("test query") # pyright: ignore[reportPrivateUsage] + + tool = AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + embedding_provider="unsupported_provider", + embedding_model="test-model", + ) + with pytest.raises(ValueError, match="Unsupported client-side embedding provider: unsupported_provider"): + await tool._get_embedding("test query") # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_abstract_base_class() -> None: + """Test abstract base class behavior.""" + with pytest.raises(NotImplementedError): + BaseAzureAISearchTool._from_config( # pyright: ignore[reportPrivateUsage] + AzureAISearchConfig( + name="test", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + ) + + +@pytest.mark.asyncio +async def test_client_initialization_errors() -> None: + """Test client initialization error handling.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + with patch("azure.search.documents.aio.SearchClient.__init__", side_effect=Exception("Connection error")): + with pytest.raises(ValueError, match="Unexpected error initializing search client: Connection error"): + await tool._get_client() # pyright: ignore[reportPrivateUsage] + + with patch( + "azure.search.documents.aio.SearchClient.__init__", side_effect=ResourceNotFoundError("Index not found") + ): + with pytest.raises(ValueError, match=f"Index '{MOCK_INDEX}' not found"): + await tool._get_client() # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_client_initialization_with_error() -> None: + """Test client initialization with various errors.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + class MockResponse: + def __init__(self, status_code: int, reason: str): + self.status_code = status_code + self.reason = reason + self.request = object() + + def text(self) -> str: + return f"{self.status_code} {self.reason}" + + mock_response = MockResponse(status_code=401, reason="Unauthorized") + with patch( + "azure.search.documents.aio.SearchClient.__init__", side_effect=HttpResponseError(response=mock_response) + ): + with pytest.raises(ValueError, match="Authentication failed"): + await tool._get_client() # pyright: ignore[reportPrivateUsage] + + mock_response = MockResponse(status_code=403, reason="Forbidden") + with patch( + "azure.search.documents.aio.SearchClient.__init__", side_effect=HttpResponseError(response=mock_response) + ): + with pytest.raises(ValueError, match="Permission denied"): + await tool._get_client() # pyright: ignore[reportPrivateUsage] + + mock_response = MockResponse(status_code=500, reason="Internal Server Error") + with patch( + "azure.search.documents.aio.SearchClient.__init__", side_effect=HttpResponseError(response=mock_response) + ): + with pytest.raises(ValueError, match="Error connecting to Azure AI Search"): + await tool._get_client() # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_search_document_processing_error(mock_search_client: AsyncMock) -> None: + """Test error handling during search document processing.""" + mock_search_client.search.return_value.__aiter__.return_value = [{"invalid": "document", "@search.score": None}] + + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + with patch.object(tool, "_get_client", return_value=mock_search_client): + results = await tool.run("test query") + assert len(results.results) == 0 + + +@pytest.mark.asyncio +async def test_search_with_expired_cache( + mock_search_client: AsyncMock, mock_search_results: List[Dict[str, Any]] +) -> None: + """Test search with expired cache.""" + mock_search_client.search.return_value.__aiter__.return_value = mock_search_results + + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + enable_caching=True, + cache_ttl_seconds=1, + ) + + with patch.object(tool, "_get_client", return_value=mock_search_client): + await tool.run("test query") + assert mock_search_client.search.call_count == 1 + + await asyncio.sleep(1.1) + + await tool.run("test query") + assert mock_search_client.search.call_count == 2 + + +@pytest.mark.asyncio +async def test_search_with_invalid_credential() -> None: + """Test search with invalid credential format.""" + with pytest.raises(ValueError, match="Invalid configuration"): + AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential={"api_key": "test-key"}, # type: ignore + ) + + +@pytest.mark.asyncio +async def test_search_with_empty_query() -> None: + """Test search with empty query.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + with pytest.raises(ValueError, match="Search query cannot be empty"): + await tool.run("") + + +@pytest.mark.asyncio +async def test_vector_search_without_query() -> None: + """Test vector search with empty query.""" + tool = AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + ) + + with pytest.raises(ValueError, match="Search query cannot be empty"): + await tool.run("") + + +@pytest.mark.asyncio +async def test_search_with_cancellation_token_already_cancelled() -> None: + """Test search with already cancelled token.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + token = CancellationToken() + token.cancel() + + with pytest.raises(asyncio.CancelledError): + await tool.run("test query", cancellation_token=token) + + +@pytest.mark.asyncio +async def test_config_validation_edge_cases() -> None: + """Test configuration validation edge cases.""" + with pytest.raises( + ValueError, + match="Invalid configuration: 1 validation error for AzureAISearchConfig\n Value error, vector_fields must be provided for vector search", + ): + AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=[], + ) + + with pytest.raises(ValueError, match="vector_fields must contain at least one field name for hybrid search"): + AzureAISearchTool.create_hybrid_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=[], + search_fields=["content"], + ) + + with pytest.raises(ValueError, match="semantic_config_name is required when query_type is 'semantic'"): + AzureAISearchTool.create_hybrid_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + search_fields=["content"], + query_type="semantic", + ) + + +@pytest.mark.asyncio +async def test_base_class_functionality() -> None: + """Test base class functionality.""" + config = AzureAISearchConfig( + name="test", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + with pytest.raises(NotImplementedError, match="BaseAzureAISearchTool.*cannot be instantiated directly"): + BaseAzureAISearchTool._from_config(config) # pyright: ignore[reportPrivateUsage] + + class TestSearchTool(BaseAzureAISearchTool): + async def _get_embedding(self, query: str) -> List[float]: + return [0.1, 0.2, 0.3] + + @classmethod + def _from_config(cls, config: AzureAISearchConfig) -> "TestSearchTool": + return cls( + name=config.name, + endpoint=config.endpoint, + index_name=config.index_name, + credential=config.credential, + ) + + tool = TestSearchTool._from_config(config) # pyright: ignore[reportPrivateUsage] + assert tool.name == "test" + assert tool.search_config.endpoint == MOCK_ENDPOINT + + +@pytest.mark.asyncio +async def test_client_cleanup_edge_cases() -> None: + """Test client cleanup edge cases.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + ) + + tool._client = None # pyright: ignore[reportPrivateUsage] + await tool.close() + + mock_client = AsyncMock() + mock_client.close.side_effect = Exception("Failed to close") + tool._client = mock_client # pyright: ignore[reportPrivateUsage] + await tool.close() + assert tool._client is None # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_token_acquisition_edge_cases() -> None: + """Test token acquisition edge cases.""" + with patch("azure.identity.DefaultAzureCredential") as mock_credential: + mock_credential.return_value.get_token.return_value = None + + tool = AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + embedding_provider="azure_openai", + embedding_model="text-embedding-ada-002", + openai_endpoint="https://test.openai.azure.com", + openai_api_version="2023-11-01", + ) + + with pytest.raises(ValueError, match="Failed to acquire token using DefaultAzureCredential for Azure OpenAI"): + await tool._get_embedding("test query") # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_hybrid_search_validation() -> None: + """Test hybrid search validation edge cases.""" + with pytest.raises(ValueError, match="semantic_config_name is required when query_type is 'semantic'"): + AzureAISearchTool.create_hybrid_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + search_fields=["content"], + query_type="semantic", + ) + + with pytest.raises(ValueError, match="openai_endpoint is required when embedding_provider is 'azure_openai'"): + AzureAISearchTool.create_hybrid_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + search_fields=["content"], + embedding_provider="azure_openai", + embedding_model="text-embedding-ada-002", + ) + + +@pytest.mark.asyncio +async def test_search_result_caching() -> None: + """Test that search results are properly cached and retrieved.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + enable_caching=True, + cache_ttl_seconds=10, + ) + + mock_results = [{"id": "1", "content": "Test", "@search.score": 0.8}] + + with patch.object(tool, "_get_client") as mock_get_client: + mock_client = AsyncMock() + mock_client.search.return_value.__aiter__.return_value = mock_results + mock_get_client.return_value = mock_client + + result1 = await tool.run("test query") + assert len(result1.results) == 1 + assert mock_client.search.call_count == 1 + + result2 = await tool.run("test query") + assert len(result2.results) == 1 + assert mock_client.search.call_count == 1 + + +@pytest.mark.asyncio +async def test_cache_expiration() -> None: + """Test that cached results expire after TTL.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + enable_caching=True, + cache_ttl_seconds=1, + ) + + mock_results = [{"id": "1", "content": "Test", "@search.score": 0.8}] + + with patch.object(tool, "_get_client") as mock_get_client: + mock_client = AsyncMock() + mock_client.search.return_value.__aiter__.return_value = mock_results + mock_get_client.return_value = mock_client + + await tool.run("test query") + assert mock_client.search.call_count == 1 + + await asyncio.sleep(1.1) + + await tool.run("test query") + assert mock_client.search.call_count == 2 + + +@pytest.mark.asyncio +async def test_search_field_validation() -> None: + """Test validation of search fields configuration.""" + tool = AzureAISearchTool.create_full_text_search( + name="test-search", endpoint=MOCK_ENDPOINT, index_name=MOCK_INDEX, credential=MOCK_CREDENTIAL, search_fields=[] + ) + assert tool.search_config.search_fields == [] + + tool = AzureAISearchTool.create_full_text_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + search_fields=["content", "content"], + ) + assert tool.search_config.search_fields == ["content", "content"] + + +@pytest.mark.asyncio +async def test_api_version_handling() -> None: + """Test handling of different API versions.""" + tool = AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + api_version="2023-11-01", + ) + assert tool.search_config.api_version == "2023-11-01" + + tool = AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + ) + assert tool.search_config.api_version == DEFAULT_API_VERSION + + with patch("autogen_ext.tools.azure._ai_search.logger") as mock_logger: + AzureAISearchTool.create_vector_search( + name="test-search", + endpoint=MOCK_ENDPOINT, + index_name=MOCK_INDEX, + credential=MOCK_CREDENTIAL, + vector_fields=["embedding"], + api_version="2023-11-01", + ) + + mock_logger.warning.assert_called_once() + warning_msg = mock_logger.warning.call_args[0][0] + assert "vector search" in warning_msg.lower() + assert "2023-11-01" in warning_msg diff --git a/python/packages/autogen-ext/tests/tools/graphrag/test_graphrag_tools.py b/python/packages/autogen-ext/tests/tools/graphrag/test_graphrag_tools.py index f77eca627a15..4bfdc993d105 100644 --- a/python/packages/autogen-ext/tests/tools/graphrag/test_graphrag_tools.py +++ b/python/packages/autogen-ext/tests/tools/graphrag/test_graphrag_tools.py @@ -9,49 +9,109 @@ from autogen_core import CancellationToken from autogen_ext.tools.graphrag import GlobalSearchTool, GlobalSearchToolReturn, LocalSearchTool, LocalSearchToolReturn from autogen_ext.tools.graphrag._config import GlobalDataConfig, LocalDataConfig -from graphrag.callbacks.llm_callbacks import BaseLLMCallback -from graphrag.model.types import TextEmbedder -from graphrag.query.llm.base import BaseLLM, BaseTextEmbedding +from graphrag.config.models.language_model_config import LanguageModelConfig +from graphrag.data_model.types import TextEmbedder from graphrag.vector_stores.base import BaseVectorStore, VectorStoreDocument, VectorStoreSearchResult -class MockLLM(BaseLLM): # type: ignore - def generate( +class MockModelOutput: + """Mock ModelOutput implementation.""" + + def __init__(self, content: str = "Mock response") -> None: + self._content = content + + @property + def content(self) -> str: + return self._content + + @property + def full_response(self) -> dict[str, Any] | None: + return {"content": self._content} + + +class MockModelResponse: + """Mock ModelResponse implementation.""" + + def __init__(self, content: str = "Mock response") -> None: + self._output = MockModelOutput(content) + self._history: list[Any] = [] + + @property + def output(self) -> MockModelOutput: + return self._output + + @property + def parsed_response(self) -> Any | None: + return None + + @property + def history(self) -> list[Any]: + return self._history + + +class MockChatModel: # type: ignore + """Mock ChatModel implementation for testing.""" + + def __init__(self) -> None: + # Create a proper LanguageModelConfig instance + self.config: LanguageModelConfig = LanguageModelConfig( + model="gpt-3.5-turbo", type="openai_chat", api_key="mock-key" + ) + + async def achat( self, - messages: str | list[Any], - streaming: bool = True, - callbacks: list[BaseLLMCallback] | None = None, + prompt: str, + history: list[Any] | None = None, **kwargs: Any, - ) -> str: - return "Mock response" + ) -> MockModelResponse: + return MockModelResponse("Mock response") - def stream_generate( - self, messages: str | list[Any], callbacks: list[BaseLLMCallback] | None = None, **kwargs: Any - ) -> Generator[str, None, None]: + async def achat_stream( + self, + prompt: str, + history: list[Any] | None = None, + **kwargs: Any, + ) -> AsyncGenerator[str, None]: yield "Mock response" - async def agenerate( + def chat( self, - messages: str | list[Any], - streaming: bool = True, - callbacks: list[BaseLLMCallback] | None = None, + prompt: str, + history: list[Any] | None = None, **kwargs: Any, - ) -> str: - return "Mock response" + ) -> MockModelResponse: + return MockModelResponse("Mock response") - async def astream_generate( # type: ignore - self, messages: str | list[Any], callbacks: list[BaseLLMCallback] | None = None, **kwargs: Any - ) -> AsyncGenerator[str, None]: + def chat_stream( + self, + prompt: str, + history: list[Any] | None = None, + **kwargs: Any, + ) -> Generator[str, None, None]: yield "Mock response" -class MockTextEmbedding(BaseTextEmbedding): # type: ignore - def embed(self, text: str, **kwargs: Any) -> list[float]: - return [0.1] * 10 +class MockEmbeddingModel: # type: ignore + """Mock EmbeddingModel implementation for testing.""" + + def __init__(self) -> None: + # Create a proper LanguageModelConfig instance + self.config: LanguageModelConfig = LanguageModelConfig( + model="text-embedding-ada-002", type="openai_embedding", api_key="mock-key" + ) + + async def aembed_batch(self, text_list: list[str], **kwargs: Any) -> list[list[float]]: + return [[0.1] * 10 for _ in text_list] async def aembed(self, text: str, **kwargs: Any) -> list[float]: return [0.1] * 10 + def embed_batch(self, text_list: list[str], **kwargs: Any) -> list[list[float]]: + return [[0.1] * 10 for _ in text_list] + + def embed(self, text: str, **kwargs: Any) -> list[float]: + return [0.1] * 10 + class MockVectorStore(BaseVectorStore): # type: ignore def __init__(self, **kwargs: Any) -> None: @@ -96,10 +156,10 @@ async def test_global_search_tool( # Create a temporary directory to simulate the data config with tempfile.TemporaryDirectory() as tempdir: # Save fixtures to parquet files - community_table = os.path.join(tempdir, "create_final_communities.parquet") - entity_table = os.path.join(tempdir, "create_final_nodes.parquet") - community_report_table = os.path.join(tempdir, "create_final_community_reports.parquet") - entity_embedding_table = os.path.join(tempdir, "create_final_entities.parquet") + community_table = os.path.join(tempdir, "communities.parquet") + entity_table = os.path.join(tempdir, "entities.parquet") + community_report_table = os.path.join(tempdir, "community_reports.parquet") + entity_embedding_table = os.path.join(tempdir, "entities.parquet") community_df_fixture.to_parquet(community_table) # type: ignore entity_df_fixture.to_parquet(entity_table) # type: ignore @@ -109,17 +169,17 @@ async def test_global_search_tool( # Initialize the data config with the temporary directory data_config = GlobalDataConfig( input_dir=tempdir, - community_table="create_final_communities", - entity_table="create_final_nodes", - community_report_table="create_final_community_reports", - entity_embedding_table="create_final_entities", + community_table="communities", + entity_table="entities", + community_report_table="community_reports", + entity_embedding_table="entities", ) # Initialize the GlobalSearchTool with mock data token_encoder = tiktoken.encoding_for_model("gpt-4o") - llm = MockLLM() + model = MockChatModel() - global_search_tool = GlobalSearchTool(token_encoder=token_encoder, llm=llm, data_config=data_config) + global_search_tool = GlobalSearchTool(token_encoder=token_encoder, model=model, data_config=data_config) with caplog.at_level("INFO"): # Example of running the tool and checking the result @@ -139,35 +199,38 @@ async def test_local_search_tool( relationship_df_fixture: pd.DataFrame, text_unit_df_fixture: pd.DataFrame, entity_embedding_fixture: pd.DataFrame, + community_df_fixture: pd.DataFrame, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, ) -> None: # Create a temporary directory to simulate the data config with tempfile.TemporaryDirectory() as tempdir: # Save fixtures to parquet files - entity_table = os.path.join(tempdir, "create_final_nodes.parquet") - relationship_table = os.path.join(tempdir, "create_final_relationships.parquet") - text_unit_table = os.path.join(tempdir, "create_final_text_units.parquet") - entity_embedding_table = os.path.join(tempdir, "create_final_entities.parquet") + entity_table = os.path.join(tempdir, "entities.parquet") + relationship_table = os.path.join(tempdir, "relationships.parquet") + text_unit_table = os.path.join(tempdir, "text_units.parquet") + entity_embedding_table = os.path.join(tempdir, "entities.parquet") + community_table = os.path.join(tempdir, "communities.parquet") entity_df_fixture.to_parquet(entity_table) # type: ignore relationship_df_fixture.to_parquet(relationship_table) # type: ignore text_unit_df_fixture.to_parquet(text_unit_table) # type: ignore entity_embedding_fixture.to_parquet(entity_embedding_table) # type: ignore + community_df_fixture.to_parquet(community_table) # type: ignore # Initialize the data config with the temporary directory data_config = LocalDataConfig( input_dir=tempdir, - entity_table="create_final_nodes", - relationship_table="create_final_relationships", - text_unit_table="create_final_text_units", - entity_embedding_table="create_final_entities", + entity_table="entities", + relationship_table="relationships", + text_unit_table="text_units", + entity_embedding_table="entities", ) # Initialize the LocalSearchTool with mock data token_encoder = tiktoken.encoding_for_model("gpt-4o") - llm = MockLLM() - embedder = MockTextEmbedding() + model = MockChatModel() + embedder = MockEmbeddingModel() # Mock the vector store def mock_vector_store_factory(*args: Any, **kwargs: dict[str, Any]) -> MockVectorStore: @@ -179,7 +242,7 @@ def mock_vector_store_factory(*args: Any, **kwargs: dict[str, Any]) -> MockVecto monkeypatch.setattr("autogen_ext.tools.graphrag._local_search.LanceDBVectorStore", mock_vector_store_factory) # type: ignore local_search_tool = LocalSearchTool( - token_encoder=token_encoder, llm=llm, embedder=embedder, data_config=data_config + token_encoder=token_encoder, model=model, embedder=embedder, data_config=data_config ) with caplog.at_level("INFO"): diff --git a/python/packages/autogen-ext/tests/tools/http/test_http_tool.py b/python/packages/autogen-ext/tests/tools/http/test_http_tool.py index 9d2898bae505..8e50e48ba926 100644 --- a/python/packages/autogen-ext/tests/tools/http/test_http_tool.py +++ b/python/packages/autogen-ext/tests/tools/http/test_http_tool.py @@ -176,7 +176,7 @@ async def test_invalid_request(test_config: ComponentModel, test_server: None) - config.config["host"] = "fake" tool = HttpTool.load_component(config) - with pytest.raises(httpx.ConnectError): + with pytest.raises((httpx.ConnectError, httpx.ConnectTimeout)): await tool.run_json({"query": "test query", "value": 42}, CancellationToken()) diff --git a/python/packages/autogen-ext/tests/tools/test_mcp_actor.py b/python/packages/autogen-ext/tests/tools/test_mcp_actor.py new file mode 100644 index 000000000000..7fc82619ce84 --- /dev/null +++ b/python/packages/autogen-ext/tests/tools/test_mcp_actor.py @@ -0,0 +1,1332 @@ +"""Tests for McpSessionActor to cover missing test coverage lines.""" + +import asyncio +import atexit +import json +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any, AsyncGenerator, Callable, Generator +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from autogen_core import Image +from autogen_core.models import ( + CreateResult, + ModelInfo, + RequestUsage, + UserMessage, +) +from autogen_ext.tools.mcp import StdioServerParams +from autogen_ext.tools.mcp._actor import ( + McpSessionActor, + _parse_sampling_content, # pyright: ignore[reportPrivateUsage] + _parse_sampling_message, # pyright: ignore[reportPrivateUsage] +) +from mcp import types as mcp_types +from mcp.shared.context import RequestContext + +# Monkey patch to prevent atexit handlers from being registered during tests +# This prevents the test suite from hanging during shutdown +original_atexit_register = atexit.register + + +def mock_atexit_register(func: Callable[[], None]) -> None: + """Mock atexit.register to prevent registration during tests.""" + pass + + +# Apply the monkey patch +atexit.register = mock_atexit_register # type: ignore[assignment] + + +@pytest.fixture +def mcp_server_params() -> StdioServerParams: + """Create server parameters that will launch the real MCP server subprocess.""" + # Get the path to the simple MCP server + server_path = Path(__file__).parent.parent / "mcp_server_comprehensive.py" + return StdioServerParams( + command="uv", + args=["run", "python", str(server_path)], + read_timeout_seconds=10, + ) + + +@pytest.fixture +def mock_model_client() -> MagicMock: + """Mock model client for testing.""" + model_client = MagicMock() + model_client.model_info = { + "vision": False, + "function_calling": False, + "json_output": False, + "family": "test-model", + "structured_output": False, + } + model_client.create = AsyncMock( + return_value=CreateResult( + content="Mock response", + finish_reason="stop", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ) + return model_client + + +@pytest.fixture +def mock_model_client_with_vision() -> MagicMock: + """Mock model client with vision support for testing.""" + model_client = MagicMock() + model_client.model_info = { + "vision": True, + "function_calling": False, + "json_output": False, + "family": "test-vision-model", + "structured_output": False, + } + model_client.create = AsyncMock( + return_value=CreateResult( + content="Mock response", + finish_reason="stop", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ) + return model_client + + +def test_parse_sampling_content_unsupported_image_without_vision() -> None: + """Test _parse_sampling_content raises error for image content when model doesn't support vision (line 56).""" + model_info: ModelInfo = { + "vision": False, + "function_calling": False, + "json_output": False, + "family": "test-model", + "structured_output": False, + } + image_content = mcp_types.ImageContent( + type="image", + data="iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==", + mimeType="image/png", + ) + + with pytest.raises(ValueError, match="Sampling model does not support image content"): + _parse_sampling_content(image_content, model_info) + + +def test_parse_sampling_content_with_vision_support() -> None: + """Test _parse_sampling_content works with vision-enabled model.""" + model_info: ModelInfo = { + "vision": True, + "function_calling": False, + "json_output": False, + "family": "test-model", + "structured_output": False, + } + image_content = mcp_types.ImageContent( + type="image", + data="iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==", + mimeType="image/png", + ) + + result = _parse_sampling_content(image_content, model_info) + assert isinstance(result, Image) + + +def test_parse_sampling_content_text() -> None: + """Test _parse_sampling_content with text content.""" + model_info: ModelInfo = { + "vision": False, + "function_calling": False, + "json_output": False, + "family": "test-model", + "structured_output": False, + } + text_content = mcp_types.TextContent(type="text", text="Hello world") + + result = _parse_sampling_content(text_content, model_info) + assert result == "Hello world" + + +def test_parse_sampling_content_unknown_type() -> None: + """Test _parse_sampling_content raises error for unknown content type.""" + model_info: ModelInfo = { + "vision": False, + "function_calling": False, + "json_output": False, + "family": "test-model", + "structured_output": False, + } + # Create a mock content with unknown type + unknown_content = MagicMock() + unknown_content.type = "unknown" + + with pytest.raises(ValueError, match="Unsupported content type"): + _parse_sampling_content(unknown_content, model_info) + + +def test_parse_sampling_message_unrecognized_role() -> None: + """Test _parse_sampling_message raises error for unrecognized role (lines 67-74).""" + model_info: ModelInfo = { + "vision": False, + "function_calling": False, + "json_output": False, + "family": "test-model", + "structured_output": False, + } + # Create a mock message with invalid role by bypassing type checking + message = MagicMock() + message.role = "system" # Invalid role that should trigger the error + message.content = mcp_types.TextContent(type="text", text="Hello") + + with pytest.raises(ValueError, match="Unrecognized message role: system"): + _parse_sampling_message(message, model_info) + + +def test_parse_sampling_message_assistant_with_non_string_content() -> None: + """Test _parse_sampling_message with assistant message containing non-string content.""" + model_info: ModelInfo = { + "vision": True, + "function_calling": False, + "json_output": False, + "family": "test-model", + "structured_output": False, + } + # Create image content for assistant message (which should fail) + image_content = mcp_types.ImageContent( + type="image", + data="iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==", + mimeType="image/png", + ) + message = mcp_types.SamplingMessage(role="assistant", content=image_content) + + # This should raise an AssertionError because assistant messages only support string content + with pytest.raises(AssertionError, match="Assistant messages only support string content"): + _parse_sampling_message(message, model_info) + + +@pytest.mark.asyncio +async def test_call_when_not_active() -> None: + """Test call method raises error when actor is not active (line 110).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + + with pytest.raises(RuntimeError, match="MCP Actor not running, call initialize\\(\\) first"): + await actor.call("list_tools") + + +@pytest.mark.asyncio +async def test_call_when_actor_task_crashed() -> None: + """Test call method raises error when actor task has crashed (line 119).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + + # Create a crashed task + async def failing_task() -> Any: + raise ValueError("Task crashed") + + actor._actor_task = asyncio.create_task(failing_task()) # type: ignore[reportPrivateUsage] # type: ignore[reportPrivateUsage] + + try: + await asyncio.sleep(0.01) # Let the task crash + actor._active = True # type: ignore[reportPrivateUsage] + + with pytest.raises(RuntimeError, match="MCP actor task crashed"): + await actor.call("list_tools") + finally: + # Clean up the task + if not actor._actor_task.done(): # type: ignore[reportPrivateUsage] + actor._actor_task.cancel() # type: ignore[reportPrivateUsage] + try: + await actor._actor_task # type: ignore[reportPrivateUsage] + except (asyncio.CancelledError, ValueError): + pass + + +@pytest.mark.asyncio +async def test_call_without_required_args() -> None: + """Test call method raises error when args are required but not provided (line 121).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._active = True # type: ignore[reportPrivateUsage] + actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] + actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage] + + with pytest.raises(ValueError, match="args is required for call_tool"): + await actor.call("call_tool") + + +@pytest.mark.asyncio +async def test_call_tool_without_name() -> None: + """Test call_tool raises error when name is not provided (line 128).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._active = True # type: ignore[reportPrivateUsage] + actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] + actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage] + + with pytest.raises(ValueError, match="name is required for call_tool"): + await actor.call("call_tool", {"name": None, "kargs": {}}) + + +@pytest.mark.asyncio +async def test_read_resource_without_uri() -> None: + """Test read_resource raises error when uri is not provided (line 132).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._active = True # type: ignore[reportPrivateUsage] + actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] + actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage] + + with pytest.raises(ValueError, match="uri is required for read_resource"): + await actor.call("read_resource", {"name": None, "kargs": {}}) + + +# @pytest.mark.asyncio +# async def test_read_resource_command_queuing() -> None: +# """Test read_resource command queuing (lines 134-137).""" +# actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) +# actor._active = True # type: ignore[reportPrivateUsage] +# actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] +# actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage] + +# # Mock the command queue to capture what gets put in it +# original_queue = actor._command_queue # type: ignore[reportPrivateUsage] +# actor._command_queue = MagicMock() # type: ignore[reportPrivateUsage] +# actor._command_queue.put = AsyncMock() # type: ignore[reportPrivateUsage] + +# # Create a task that will be cancelled to avoid hanging +# call_task = asyncio.create_task(actor.call("read_resource", {"name": None, "kargs": {"uri": "file:///test.txt"}})) + +# # Give it a brief moment to queue the command +# await asyncio.sleep(0.001) + +# # Cancel the task to avoid hanging +# call_task.cancel() + +# # Wait for the cancellation to complete +# try: +# await call_task +# except asyncio.CancelledError: +# pass # Expected + +# # Verify the command was queued correctly (this covers lines 134-137) +# actor._command_queue.put.assert_called_once() # type: ignore[reportPrivateUsage] +# call_args = actor._command_queue.put.call_args[0][0] # type: ignore[reportPrivateUsage] +# assert call_args["type"] == "read_resource" +# assert call_args["uri"] == "file:///test.txt" +# assert "future" in call_args + +# # Restore original queue +# actor._command_queue = original_queue # type: ignore[reportPrivateUsage] + + +# @pytest.mark.asyncio +# async def test_get_prompt_without_name() -> None: +# """Test get_prompt raises error when name is not provided (lines 139-142).""" +# actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) +# actor._active = True # type: ignore[reportPrivateUsage] +# actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] +# actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage] + +# with pytest.raises(ValueError, match="name is required for get_prompt"): +# await actor.call("get_prompt", {"name": None, "kargs": {}}) + + +# @pytest.mark.asyncio +# async def test_get_prompt_command_queuing() -> None: +# """Test get_prompt command queuing (lines 139-142).""" +# actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) +# actor._active = True # type: ignore[reportPrivateUsage] +# actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] +# actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage] + +# # Mock the command queue to capture what gets put in it +# original_queue = actor._command_queue # type: ignore[reportPrivateUsage] +# actor._command_queue = MagicMock() # type: ignore[reportPrivateUsage] +# actor._command_queue.put = AsyncMock() # type: ignore[reportPrivateUsage] + +# # This will fail when trying to await the future, but we're testing the command queuing logic +# try: +# await actor.call("get_prompt", {"name": "test_prompt", "kargs": {"arguments": {"arg1": "value1"}}}) +# except Exception: +# pass # Expected to fail, we're just testing the queuing logic + +# # Verify the command was queued correctly (this covers lines 139-142) +# actor._command_queue.put.assert_called_once() # type: ignore[reportPrivateUsage] +# call_args = actor._command_queue.put.call_args[0][0] # type: ignore[reportPrivateUsage] +# assert call_args["type"] == "get_prompt" +# assert call_args["name"] == "test_prompt" +# assert call_args["args"] == {"arg1": "value1"} +# assert "future" in call_args + +# # Restore original queue +# actor._command_queue = original_queue # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_call_unknown_command_type() -> None: + """Test call method raises error for unknown command type (line 147).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._active = True # type: ignore[reportPrivateUsage] + actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] + actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage] + + with pytest.raises(ValueError, match="Unknown command type: unknown_command"): + await actor.call("unknown_command") + + +@pytest.mark.asyncio +async def test_close_when_not_active() -> None: + """Test close method early return when not active (line 152).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._active = False # type: ignore[reportPrivateUsage] + actor._actor_task = None # type: ignore[reportPrivateUsage] + + # This should return early without doing anything + await actor.close() + assert actor._shutdown_future is None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_sampling_callback_without_model_client() -> None: + """Test sampling callback returns error when no model client is available (line 194).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]), model_client=None) + + mock_context = MagicMock(spec=RequestContext) + params = mcp_types.CreateMessageRequestParams( + messages=[mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello"))], + maxTokens=100, + ) + + result = await actor._sampling_callback(mock_context, params) # type: ignore[reportPrivateUsage] + + assert isinstance(result, mcp_types.ErrorData) + assert result.code == mcp_types.INVALID_REQUEST + assert "No model client available" in result.message + + +@pytest.mark.asyncio +async def test_sampling_callback_message_processing_error() -> None: + """Test sampling callback handles message processing errors (lines 226-227, 229-233).""" + # Create a model client for the actor + model_client = MagicMock() + model_client.model_info = { + "vision": False, + "function_calling": False, + "json_output": False, + "family": "test-model", + "structured_output": False, + } + + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]), model_client=model_client) + + mock_context = MagicMock(spec=RequestContext) + + # Create a valid SamplingMessage but with invalid role that will cause parsing error + # We'll patch the message after creation to bypass Pydantic validation + valid_message = mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello")) + # Now change the role to something invalid that will cause _parse_sampling_message to fail + valid_message.role = "invalid_role" # type: ignore + + params = mcp_types.CreateMessageRequestParams( + messages=[valid_message], + maxTokens=100, + ) + + result = await actor._sampling_callback(mock_context, params) # type: ignore[reportPrivateUsage] + + assert isinstance(result, mcp_types.ErrorData) + assert result.code == mcp_types.INVALID_PARAMS + assert "Error processing sampling messages" in result.message + + +@pytest.mark.asyncio +async def test_sampling_callback_model_client_error() -> None: + """Test sampling callback handles model client errors (lines 235-239).""" + failing_model_client = MagicMock() + failing_model_client.model_info = {"vision": False, "family": "test-model"} + failing_model_client.create = AsyncMock(side_effect=Exception("Model API error")) + + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]), model_client=failing_model_client) + + mock_context = MagicMock(spec=RequestContext) + params = mcp_types.CreateMessageRequestParams( + messages=[mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello"))], + maxTokens=100, + ) + + result = await actor._sampling_callback(mock_context, params) # type: ignore[reportPrivateUsage] + + assert isinstance(result, mcp_types.ErrorData) + assert result.code == mcp_types.INTERNAL_ERROR + assert "Error sampling from model client" in result.message + assert "Model API error" in str(result.data) + + +@pytest.mark.asyncio +async def test_run_actor_exception_handling() -> None: + """Test _run_actor exception handling for various command types (lines 244-268).""" + # This test focuses on covering the exception handling code paths in _run_actor + # We'll test this by creating a simplified scenario that exercises those paths + + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + + # Create a mock session that will raise exceptions by modifying our mock session + @asynccontextmanager + async def mock_failing_session( + server_params: Any, sampling_callback: Any = None + ) -> AsyncGenerator[MagicMock, None]: + mock_session = MagicMock() + mock_session.initialize = AsyncMock( + return_value=mcp_types.InitializeResult( + protocolVersion="1.0", + capabilities=mcp_types.ServerCapabilities(), + serverInfo=mcp_types.Implementation(name="test", version="1.0"), + ) + ) + # Make sure all the session methods raise exceptions to cover the exception handling + mock_session.call_tool = MagicMock(side_effect=Exception("Tool error")) + mock_session.list_tools = MagicMock(side_effect=Exception("List tools error")) + yield mock_session + + with patch("autogen_ext.tools.mcp._actor.create_mcp_server_session", mock_failing_session): # type: ignore[reportPrivateUsage] + # Start the actor task + actor._active = True # type: ignore[reportPrivateUsage] + actor_task = asyncio.create_task(actor._run_actor()) # type: ignore[reportPrivateUsage] + + try: + # Give it a moment to initialize + await asyncio.sleep(0.05) + + # Test one command that will trigger exception handling (covers lines 244-268) + future: asyncio.Future[Any] = asyncio.Future() + cmd: dict[str, Any] = {"type": "call_tool", "name": "test_tool", "args": {}, "future": future} + await actor._command_queue.put(cmd) # type: ignore[reportPrivateUsage] + + # Wait a bit for command to be processed + await asyncio.sleep(0.05) + + # Send shutdown command + shutdown_future: asyncio.Future[Any] = asyncio.Future() + await actor._command_queue.put({"type": "shutdown", "future": shutdown_future}) # type: ignore[reportPrivateUsage] + + # Wait for actor to finish + try: + await asyncio.wait_for(actor_task, timeout=1.0) + except asyncio.TimeoutError: + pass # Expected if task doesn't finish properly + + # The key test: verify that the future was set with an exception + # This proves the exception handling code (lines 244-268) was executed + assert future.done() + assert future.exception() is not None + assert "Tool error" in str(future.exception()) + finally: + # Ensure the task is cancelled and cleaned up + if not actor_task.done(): + actor_task.cancel() + try: + await actor_task + except asyncio.CancelledError: + pass + + +@pytest.mark.asyncio +async def test_run_actor_session_exception() -> None: + """Test _run_actor handles session creation exceptions (lines 274-288).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + + # Mock session creation to raise an exception + with patch("autogen_ext.tools.mcp._actor.create_mcp_server_session", side_effect=Exception("Session error")): # type: ignore[reportPrivateUsage] + actor._active = True # type: ignore[reportPrivateUsage] + actor_task = asyncio.create_task(actor._run_actor()) # type: ignore[reportPrivateUsage] + + try: + # Wait for the task to complete + await asyncio.wait_for(actor_task, timeout=1.0) + except asyncio.TimeoutError: + # If it doesn't complete, cancel it + actor_task.cancel() + try: + await actor_task + except asyncio.CancelledError: + pass + + # Check that the actor is no longer active + assert not actor._active # type: ignore[reportPrivateUsage] + assert actor._actor_task is None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_run_actor_shutdown_future_exception() -> None: + """Test _run_actor sets exception on shutdown future when session fails.""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._shutdown_future = asyncio.Future() # type: ignore[reportPrivateUsage] + + # Mock session creation to raise an exception + with patch("autogen_ext.tools.mcp._actor.create_mcp_server_session", side_effect=Exception("Session error")): # type: ignore[reportPrivateUsage] + actor._active = True # type: ignore[reportPrivateUsage] + actor_task = asyncio.create_task(actor._run_actor()) # type: ignore[reportPrivateUsage] + + try: + # Wait for the task to complete + await asyncio.wait_for(actor_task, timeout=1.0) + except asyncio.TimeoutError: + # If it doesn't complete, cancel it + actor_task.cancel() + try: + await actor_task + except asyncio.CancelledError: + pass + + # Check that shutdown future has the exception + assert actor._shutdown_future.done() # type: ignore[reportPrivateUsage] + assert actor._shutdown_future.exception() is not None # type: ignore[reportPrivateUsage] + + +def test_sync_shutdown_when_not_active() -> None: + """Test _sync_shutdown early return when not active (line 297).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._active = False # type: ignore[reportPrivateUsage] + actor._actor_task = None # type: ignore[reportPrivateUsage] + + # This should return early without doing anything + actor._sync_shutdown() # type: ignore[reportPrivateUsage] + + +def test_sync_shutdown_no_event_loop() -> None: + """Test _sync_shutdown handles RuntimeError when no event loop (line 310).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._active = True # type: ignore[reportPrivateUsage] + actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] + + # Mock get_event_loop to raise RuntimeError + with patch("asyncio.get_event_loop", side_effect=RuntimeError("No event loop")): + # This should return early due to the RuntimeError + actor._sync_shutdown() # type: ignore[reportPrivateUsage] + + +def test_sync_shutdown_closed_loop() -> None: + """Test _sync_shutdown handles closed event loop.""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._active = True # type: ignore[reportPrivateUsage] + actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] + + # Mock event loop that is closed + mock_loop = MagicMock() + mock_loop.is_closed.return_value = True + + with patch("asyncio.get_event_loop", return_value=mock_loop): + # This should return early due to the closed loop + actor._sync_shutdown() # type: ignore[reportPrivateUsage] + + +def test_sync_shutdown_running_loop() -> None: + """Test _sync_shutdown creates task when loop is running.""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._active = True # type: ignore[reportPrivateUsage] + actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] + + # Mock event loop that is running + mock_loop = MagicMock() + mock_loop.is_closed.return_value = False + mock_loop.is_running.return_value = True + mock_loop.create_task = MagicMock() + + with patch("asyncio.get_event_loop", return_value=mock_loop): + actor._sync_shutdown() # type: ignore[reportPrivateUsage] + # Should create a task to close the actor + mock_loop.create_task.assert_called_once() + + +def test_sync_shutdown_non_running_loop() -> None: + """Test _sync_shutdown runs until complete when loop is not running.""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._active = True # type: ignore[reportPrivateUsage] + actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] + + # Mock event loop that is not running + mock_loop = MagicMock() + mock_loop.is_closed.return_value = False + mock_loop.is_running.return_value = False + mock_loop.run_until_complete = MagicMock() + + with patch("asyncio.get_event_loop", return_value=mock_loop): + actor._sync_shutdown() # type: ignore[reportPrivateUsage] + # Should run until complete + mock_loop.run_until_complete.assert_called_once() + + +def test_to_config() -> None: + """Test _to_config method.""" + server_params = StdioServerParams(command="echo", args=["test"]) + actor = McpSessionActor(server_params) + + config = actor._to_config() # type: ignore[reportPrivateUsage] + assert config.server_params == server_params + + +def test_from_config() -> None: + """Test _from_config class method.""" + from autogen_ext.tools.mcp._actor import McpSessionActorConfig # type: ignore[reportPrivateUsage] + + server_params = StdioServerParams(command="echo", args=["test"]) + config = McpSessionActorConfig(server_params=server_params) + + actor = McpSessionActor._from_config(config) # type: ignore[reportPrivateUsage] + assert actor.server_params == server_params + + +@pytest.mark.asyncio +async def test_initialize_result_property() -> None: + """Test initialize_result property.""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + + # Initially should be None + assert actor.initialize_result is None + + # Set a mock result + mock_result = mcp_types.InitializeResult( + protocolVersion="1.0", + capabilities=mcp_types.ServerCapabilities(), + serverInfo=mcp_types.Implementation(name="test", version="1.0"), + ) + actor._initialize_result = mock_result # type: ignore[reportPrivateUsage] + + assert actor.initialize_result == mock_result + + +@pytest.mark.asyncio +async def test_actor_initialization() -> None: + """Test actor initialization sets up correctly.""" + server_params = StdioServerParams(command="echo", args=["test"]) + model_client = MagicMock() + + actor = McpSessionActor(server_params, model_client=model_client) + + # Check initial state + assert actor.server_params == server_params + assert actor._model_client == model_client # type: ignore[reportPrivateUsage] + assert actor.name == "mcp_session_actor" + assert actor.description == "MCP session actor" + assert not actor._active # type: ignore[reportPrivateUsage] + assert actor._actor_task is None # type: ignore[reportPrivateUsage] + assert actor._shutdown_future is None # type: ignore[reportPrivateUsage] + assert actor._initialize_result is None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_initialize_method(mcp_server_params: Any) -> None: + """Test initialize method.""" + actor = McpSessionActor(mcp_server_params) + + await actor.initialize() + + assert actor._active # type: ignore[reportPrivateUsage] + assert actor._actor_task is not None # type: ignore[reportPrivateUsage] + + # Clean up + await actor.close() + + +@pytest.mark.asyncio +async def test_call_with_valid_list_commands() -> None: + """Test call method with valid list commands.""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._active = True # type: ignore[reportPrivateUsage] + actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] + actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage] + + # Mock the command queue to capture what gets put in it + original_queue = actor._command_queue # type: ignore[reportPrivateUsage] + actor._command_queue = MagicMock() # type: ignore[reportPrivateUsage] + actor._command_queue.put = AsyncMock() # type: ignore[reportPrivateUsage] + + # Test all valid list commands - we don't await the result, just test the queuing + test_commands = ["list_tools", "list_prompts", "list_resources", "list_resource_templates", "shutdown"] + created_tasks: list[asyncio.Task[Any]] = [] + + try: + for cmd_type in test_commands: + # Create a task but don't await it (to avoid hanging) + call_task = asyncio.create_task(actor.call(cmd_type)) + created_tasks.append(call_task) + + # Give it a brief moment to queue the command + await asyncio.sleep(0.001) + + # Verify the command was queued correctly + actor._command_queue.put.assert_called() # type: ignore[reportPrivateUsage] + call_args = actor._command_queue.put.call_args[0][0] # type: ignore[reportPrivateUsage] + assert call_args["type"] == cmd_type + assert "future" in call_args + actor._command_queue.put.reset_mock() # type: ignore[reportPrivateUsage] + + finally: + # Clean up all created tasks + for task in created_tasks: + if not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + # Restore original queue + actor._command_queue = original_queue # type: ignore[reportPrivateUsage] + + +# Integration tests using the real MCP server +@pytest.mark.asyncio +async def test_actor_basic_functionality(mcp_server_params: Any) -> None: + """Test basic actor functionality with real MCP server.""" + actor = McpSessionActor(mcp_server_params) + + try: + # Initialize the actor + await actor.initialize() + assert actor._active # type: ignore[reportPrivateUsage] + assert actor._actor_task is not None # type: ignore[reportPrivateUsage] + + # Test listing tools + tools_future = await actor.call("list_tools") + tools_result: mcp_types.ListToolsResult = await tools_future # type: ignore + assert len(tools_result.tools) == 2 # echo and get_time + tool_names = [tool.name for tool in tools_result.tools] + assert "echo" in tool_names + assert "get_time" in tool_names + + # Test calling a tool + call_future = await actor.call("call_tool", {"name": "echo", "kargs": {"text": "Hello World"}}) + call_result: mcp_types.CallToolResult = await call_future # type: ignore + assert call_result.content[0].text == "Echo: Hello World" # type: ignore + + finally: + await actor.close() + + +@pytest.mark.asyncio +async def test_actor_prompt_operations(mcp_server_params: Any) -> None: + """Test actor prompt operations with real MCP server.""" + actor = McpSessionActor(mcp_server_params) + + try: + await actor.initialize() + await asyncio.sleep(0.1) + + # Test listing prompts + prompts_future = await actor.call("list_prompts") + prompts_result: mcp_types.ListPromptsResult = await prompts_future # type: ignore + assert len(prompts_result.prompts) == 2 # code_review and documentation + prompt_names = [prompt.name for prompt in prompts_result.prompts] + assert "code_review" in prompt_names + assert "documentation" in prompt_names + + # Test getting a prompt with arguments + prompt_future = await actor.call( + "get_prompt", + {"name": "code_review", "kargs": {"arguments": {"code": "print('hello')", "language": "python"}}}, + ) + prompt_result: mcp_types.GetPromptResult = await prompt_future # type: ignore + assert prompt_result.description is not None and "python" in prompt_result.description + assert "print('hello')" in prompt_result.messages[0].content.text # type: ignore + + finally: + await actor.close() + + +@pytest.mark.asyncio +async def test_actor_resource_operations(mcp_server_params: Any) -> None: + """Test actor resource operations with real MCP server.""" + actor = McpSessionActor(mcp_server_params) + + try: + await actor.initialize() + await asyncio.sleep(0.1) + + # Test listing resources + resources_future = await actor.call("list_resources") + resources_result: mcp_types.ListResourcesResult = await resources_future # type: ignore + assert len(resources_result.resources) == 2 # users and projects + resource_names = [resource.name for resource in resources_result.resources] + assert "Company Users" in resource_names + assert "Active Projects" in resource_names + + # Test reading a resource + read_future = await actor.call("read_resource", {"name": None, "kargs": {"uri": "file:///company/users.json"}}) + read_result: mcp_types.ReadResourceResult = await read_future # type: ignore + users_data = json.loads(read_result.contents[0].text) # type: ignore + assert isinstance(users_data, list) + assert len(users_data) == 3 # type: ignore[reportUnknownArgumentType] + assert users_data[0]["name"] == "Alice" + + finally: + await actor.close() + + +@pytest.mark.asyncio +async def test_actor_tool_failure_handling(mcp_server_params: Any) -> None: + """Test actor handles tool failures correctly.""" + actor = McpSessionActor(mcp_server_params) + + try: + await actor.initialize() + await asyncio.sleep(0.1) + + # Test calling an unknown tool - the server should return an error result + call_future = await actor.call("call_tool", {"name": "unknown_tool", "kargs": {}}) + call_result: mcp_types.CallToolResult = await call_future # type: ignore + # The server returns an error but doesn't raise an exception + assert call_result.isError is True # type: ignore + assert "Unknown tool" in call_result.content[0].text # type: ignore + + finally: + await actor.close() + + +@pytest.mark.asyncio +async def test_actor_with_model_client_sampling(mcp_server_params: Any, mock_model_client: Any) -> None: + """Test actor with model client for sampling operations.""" + actor = McpSessionActor(mcp_server_params, model_client=mock_model_client) + + try: + await actor.initialize() + await asyncio.sleep(0.1) + + # Test sampling callback functionality + mock_context = MagicMock(spec=RequestContext) + params = mcp_types.CreateMessageRequestParams( + messages=[ + mcp_types.SamplingMessage( + role="user", content=mcp_types.TextContent(type="text", text="Hello from test") + ) + ], + maxTokens=100, + ) + + result = await actor._sampling_callback(mock_context, params) # type: ignore[reportPrivateUsage] + + assert isinstance(result, mcp_types.CreateMessageResult) + assert result.role == "assistant" + assert isinstance(result.content, mcp_types.TextContent) + assert result.content.text == "Mock response" + assert result.model == "test-model" + + finally: + await actor.close() + + +# Integration tests with real MCP server +@pytest.mark.asyncio +async def test_actor(mcp_server_params: Any) -> None: + """Test actor with real MCP server subprocess.""" + actor = McpSessionActor(mcp_server_params) + + try: + # Initialize the actor + await actor.initialize() + assert actor._active # type: ignore[reportPrivateUsage] + assert actor._actor_task is not None # type: ignore[reportPrivateUsage] + + # Test listing tools + tools_future = await actor.call("list_tools") + tools_result: mcp_types.ListToolsResult = await tools_future # type: ignore + assert len(tools_result.tools) == 2 # echo and get_time + tool_names = [tool.name for tool in tools_result.tools] + assert "echo" in tool_names + assert "get_time" in tool_names + + # Test calling the echo tool + call_future = await actor.call("call_tool", {"name": "echo", "kargs": {"text": "Hello World"}}) + call_result: mcp_types.CallToolResult = await call_future # type: ignore + assert call_result.content[0].text == "Echo: Hello World" # type: ignore + + # Test calling the get_time tool + time_future = await actor.call("call_tool", {"name": "get_time", "kargs": {}}) + time_result: mcp_types.CallToolResult = await time_future # type: ignore + assert "Current time:" in time_result.content[0].text # type: ignore + + finally: + await actor.close() + + +@pytest.mark.asyncio +async def test_actor_prompts(mcp_server_params: Any) -> None: + """Test actor prompt operations with real MCP server.""" + actor = McpSessionActor(mcp_server_params) + + try: + await actor.initialize() + + # Test listing prompts + prompts_future = await actor.call("list_prompts") + prompts_result: mcp_types.ListPromptsResult = await prompts_future # type: ignore + assert len(prompts_result.prompts) == 2 # code_review and documentation + prompt_names = [prompt.name for prompt in prompts_result.prompts] + assert "code_review" in prompt_names + assert "documentation" in prompt_names + + # Test getting a prompt + prompt_future = await actor.call( + "get_prompt", + {"name": "code_review", "kargs": {"arguments": {"code": "print('hello')", "language": "python"}}}, + ) + prompt_result: mcp_types.GetPromptResult = await prompt_future # type: ignore + assert prompt_result.description is not None and "python" in prompt_result.description + assert "print('hello')" in prompt_result.messages[0].content.text # type: ignore + + finally: + await actor.close() + + +@pytest.mark.asyncio +async def test_actor_resources(mcp_server_params: Any) -> None: + """Test actor resource operations with real MCP server.""" + actor = McpSessionActor(mcp_server_params) + + try: + await actor.initialize() + + # Test listing resources + resources_future = await actor.call("list_resources") + resources_result: mcp_types.ListResourcesResult = await resources_future # type: ignore + assert len(resources_result.resources) == 2 # users and projects + resource_names = [resource.name for resource in resources_result.resources] + assert "Company Users" in resource_names + assert "Active Projects" in resource_names + + # Test reading a resource + read_future = await actor.call("read_resource", {"name": None, "kargs": {"uri": "file:///company/users.json"}}) + read_result: mcp_types.ReadResourceResult = await read_future # type: ignore + # The real server returns content in the ReadResourceResult + users_data = json.loads(read_result.contents[0].text) # type: ignore + assert isinstance(users_data, list) + assert len(users_data) == 3 # type: ignore[reportUnknownArgumentType] + assert users_data[0]["name"] == "Alice" + + finally: + await actor.close() + + +@pytest.mark.asyncio +async def test_actor_unknown_tool(mcp_server_params: Any) -> None: + """Test actor handles unknown tools with real MCP server.""" + actor = McpSessionActor(mcp_server_params) + + try: + await actor.initialize() + + # Test calling an unknown tool - the server should return an error result + call_future = await actor.call("call_tool", {"name": "unknown_tool", "kargs": {}}) + call_result: mcp_types.CallToolResult = await call_future # type: ignore + # The server returns an error but doesn't raise an exception + assert call_result.isError is True # type: ignore + assert "Unknown tool" in call_result.content[0].text # type: ignore + + finally: + await actor.close() + + +@pytest.fixture +def clean_actor() -> Generator[Callable[..., McpSessionActor], None, None]: + """Fixture to track and clean up actors created in tests.""" + actors: list[McpSessionActor] = [] + + def create_actor(*args: Any, **kwargs: Any) -> McpSessionActor: + actor = McpSessionActor(*args, **kwargs) + actors.append(actor) + return actor + + yield create_actor + + # Clean up all actors + for actor in actors: + if hasattr(actor, "_active") and actor._active: # type: ignore[reportPrivateUsage] + try: + # Try to close the actor properly + loop = asyncio.get_event_loop() + if loop.is_running(): + loop.create_task(actor.close()) + else: + loop.run_until_complete(actor.close()) + except Exception: + # If we can't close it properly, at least deactivate it + actor._active = False # type: ignore[reportPrivateUsage] + if hasattr(actor, "_actor_task") and actor._actor_task: # type: ignore[reportPrivateUsage] + actor._actor_task.cancel() # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_sampling_callback_with_system_prompt() -> None: + """Test sampling callback with systemPrompt parameter (line 177).""" + model_client = MagicMock() + model_client.model_info = { + "vision": False, + "function_calling": False, + "json_output": False, + "family": "test-model", + "structured_output": False, + } + model_client.create = AsyncMock( + return_value=CreateResult( + content="Mock response", + finish_reason="stop", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ) + + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]), model_client=model_client) + + mock_context = MagicMock(spec=RequestContext) + params = mcp_types.CreateMessageRequestParams( + messages=[mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello"))], + maxTokens=100, + systemPrompt="You are a helpful assistant.", + ) + + result = await actor._sampling_callback(mock_context, params) # type: ignore[reportPrivateUsage] + + assert isinstance(result, mcp_types.CreateMessageResult) + assert result.role == "assistant" + assert isinstance(result.content, mcp_types.TextContent) + assert result.content.text == "Mock response" + + # Verify that the model client was called with the system prompt + model_client.create.assert_called_once() + call_args = model_client.create.call_args[1] + messages = call_args["messages"] + assert len(messages) == 2 # SystemMessage + UserMessage + assert messages[0].content == "You are a helpful assistant." + + +@pytest.mark.asyncio +async def test_sampling_callback_with_non_string_content() -> None: + """Test sampling callback when model returns non-string content (line 194).""" + model_client = MagicMock() + model_client.model_info = { + "vision": False, + "function_calling": False, + "json_output": False, + "family": "test-model", + "structured_output": False, + } + # Mock the model to return a non-string content + non_string_content = {"data": "complex object"} + model_client.create = AsyncMock( + return_value=CreateResult( + content=str(non_string_content), # Convert to string for valid content + finish_reason="stop", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ) + + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]), model_client=model_client) + + mock_context = MagicMock(spec=RequestContext) + params = mcp_types.CreateMessageRequestParams( + messages=[mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello"))], + maxTokens=100, + ) + + result = await actor._sampling_callback(mock_context, params) # type: ignore[reportPrivateUsage] + + assert isinstance(result, mcp_types.CreateMessageResult) + assert result.role == "assistant" + assert isinstance(result.content, mcp_types.TextContent) + # Should be converted to string + assert result.content.text == "{'data': 'complex object'}" + + +@pytest.mark.asyncio +async def test_run_actor_all_command_types_exception_handling() -> None: + """Test _run_actor exception handling for all command types (lines 232-233, 238-239, 244-245, 250-251, 256-263).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + + # Create a mock session that will raise exceptions for all command types + @asynccontextmanager + async def mock_failing_session( + server_params: Any, sampling_callback: Any = None + ) -> AsyncGenerator[MagicMock, None]: + mock_session = MagicMock() + mock_session.initialize = AsyncMock( + return_value=mcp_types.InitializeResult( + protocolVersion="1.0", + capabilities=mcp_types.ServerCapabilities(), + serverInfo=mcp_types.Implementation(name="test", version="1.0"), + ) + ) + # Make all session methods raise exceptions + mock_session.call_tool = MagicMock(side_effect=Exception("call_tool error")) + mock_session.read_resource = MagicMock(side_effect=Exception("read_resource error")) + mock_session.get_prompt = MagicMock(side_effect=Exception("get_prompt error")) + mock_session.list_tools = MagicMock(side_effect=Exception("list_tools error")) + mock_session.list_prompts = MagicMock(side_effect=Exception("list_prompts error")) + mock_session.list_resources = MagicMock(side_effect=Exception("list_resources error")) + mock_session.list_resource_templates = MagicMock(side_effect=Exception("list_resource_templates error")) + yield mock_session + + with patch("autogen_ext.tools.mcp._actor.create_mcp_server_session", mock_failing_session): # type: ignore[reportPrivateUsage] + # Start the actor task + actor._active = True # type: ignore[reportPrivateUsage] + actor_task = asyncio.create_task(actor._run_actor()) # type: ignore[reportPrivateUsage] + + try: + # Give it a moment to initialize + await asyncio.sleep(0.05) + + # Test all command types that can raise exceptions + commands_to_test: list[dict[str, Any]] = [ + {"type": "call_tool", "name": "test_tool", "args": {}}, + {"type": "read_resource", "uri": "test://resource"}, + {"type": "get_prompt", "name": "test_prompt", "args": {}}, + {"type": "list_tools"}, + {"type": "list_prompts"}, + {"type": "list_resources"}, + {"type": "list_resource_templates"}, + ] + + futures: list[asyncio.Future[Any]] = [] + for cmd in commands_to_test: + future: asyncio.Future[Any] = asyncio.Future() + cmd["future"] = future + await actor._command_queue.put(cmd) # type: ignore[reportPrivateUsage] + futures.append(future) + + # Wait a bit for commands to be processed + await asyncio.sleep(0.1) + + # Send shutdown command + shutdown_future: asyncio.Future[Any] = asyncio.Future() + await actor._command_queue.put({"type": "shutdown", "future": shutdown_future}) # type: ignore[reportPrivateUsage] + + # Wait for actor to finish + try: + await asyncio.wait_for(actor_task, timeout=1.0) + except asyncio.TimeoutError: + pass # Expected if task doesn't finish properly + + # Verify that all futures were set with exceptions + for i, future in enumerate(futures): + assert future.done(), f"Future {i} was not completed" + assert future.exception() is not None, f"Future {i} should have an exception" + + finally: + # Ensure the task is cancelled and cleaned up + if not actor_task.done(): + actor_task.cancel() + try: + await actor_task + except asyncio.CancelledError: + pass + + +@pytest.mark.asyncio +async def test_close_with_shutdown_await() -> None: + """Test close method waits for shutdown future (line 140).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._active = True # type: ignore[reportPrivateUsage] + actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] + actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage] + + # Mock the command queue and actor task + original_queue = actor._command_queue # type: ignore[reportPrivateUsage] + actor._command_queue = MagicMock() # type: ignore[reportPrivateUsage] + actor._command_queue.put = AsyncMock() # type: ignore[reportPrivateUsage] + + # Create a shutdown future that will be set + shutdown_future: asyncio.Future[Any] = asyncio.Future() + + # Mock the close method to set the shutdown future after a delay + async def mock_close() -> None: + actor._shutdown_future = shutdown_future # type: ignore[reportPrivateUsage] + await actor._command_queue.put({"type": "shutdown", "future": shutdown_future}) # type: ignore[reportPrivateUsage] + # Simulate the shutdown completion + await asyncio.sleep(0.01) + shutdown_future.set_result("ok") + + # Replace the close method temporarily + original_close = actor.close + actor.close = mock_close # type: ignore[method-assign] + + try: + # This should complete without hanging + await actor.close() + assert shutdown_future.done() + assert shutdown_future.result() == "ok" + finally: + # Restore original queue and close method + actor._command_queue = original_queue # type: ignore[reportPrivateUsage] + actor.close = original_close # type: ignore[method-assign] + + +def test_parse_sampling_message_user_role() -> None: + """Test _parse_sampling_message with user role (line 69).""" + model_info: ModelInfo = { + "vision": False, + "function_calling": False, + "json_output": False, + "family": "test-model", + "structured_output": False, + } + message = mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello user")) + + result = _parse_sampling_message(message, model_info) + + assert isinstance(result, UserMessage) + assert result.content == ["Hello user"] + + +@pytest.mark.asyncio +async def test_call_tool_command_queuing() -> None: + """Test call_tool command queuing (line 140).""" + actor = McpSessionActor(StdioServerParams(command="echo", args=["test"])) + actor._active = True # type: ignore[reportPrivateUsage] + actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage] + actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage] + + # Mock the command queue to capture what gets put in it + original_queue = actor._command_queue # type: ignore[reportPrivateUsage] + actor._command_queue = MagicMock() # type: ignore[reportPrivateUsage] + actor._command_queue.put = AsyncMock() # type: ignore[reportPrivateUsage] + + # Create a task that will be cancelled to avoid hanging + call_task = asyncio.create_task(actor.call("call_tool", {"name": "test_tool", "kargs": {"param": "value"}})) + + # Give it a brief moment to queue the command + await asyncio.sleep(0.001) + + # Cancel the task to avoid hanging + call_task.cancel() + + # Wait for the cancellation to complete + try: + await call_task + except asyncio.CancelledError: + pass # Expected + + # Verify the command was queued correctly (this covers line 140) + actor._command_queue.put.assert_called_once() # type: ignore[reportPrivateUsage] + call_args = actor._command_queue.put.call_args[0][0] # type: ignore[reportPrivateUsage] + assert call_args["type"] == "call_tool" + assert call_args["name"] == "test_tool" + assert call_args["args"] == {"param": "value"} + assert "future" in call_args + + # Restore original queue + actor._command_queue = original_queue # type: ignore[reportPrivateUsage] + + +def test_parse_sampling_message_user_role_branch() -> None: + """Test _parse_sampling_message with user role (line 69).""" + model_info: ModelInfo = { + "vision": False, + "function_calling": False, + "json_output": False, + "family": "test-model", + "structured_output": False, + } + message = mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello user test")) + + result = _parse_sampling_message(message, model_info) + + # This specifically tests the user role branch (line 69) + assert isinstance(result, UserMessage) + assert result.content == ["Hello user test"] diff --git a/python/packages/autogen-ext/tests/tools/test_mcp_tools.py b/python/packages/autogen-ext/tests/tools/test_mcp_tools.py index 215521aa9529..fbcbccc3e6bb 100644 --- a/python/packages/autogen-ext/tests/tools/test_mcp_tools.py +++ b/python/packages/autogen-ext/tests/tools/test_mcp_tools.py @@ -1,16 +1,37 @@ +import asyncio import logging +import os +import threading +from typing import cast from unittest.mock import AsyncMock, MagicMock import pytest +from _pytest.logging import LogCaptureFixture # type: ignore[import] from autogen_core import CancellationToken +from autogen_core.tools import Workbench +from autogen_core.utils import schema_to_pydantic_model from autogen_ext.tools.mcp import ( + McpSessionActor, + McpWorkbench, SseMcpToolAdapter, SseServerParams, StdioMcpToolAdapter, StdioServerParams, + StreamableHttpMcpToolAdapter, + StreamableHttpServerParams, + create_mcp_server_session, + mcp_server_tools, ) -from json_schema_to_pydantic import create_model from mcp import ClientSession, Tool +from mcp.types import ( + Annotations, + EmbeddedResource, + ImageContent, + ResourceLink, + TextContent, + TextResourceContents, +) +from pydantic.networks import AnyUrl @pytest.fixture @@ -44,6 +65,19 @@ def sample_sse_tool() -> Tool: ) +@pytest.fixture +def sample_streamable_http_tool() -> Tool: + return Tool( + name="test_streamable_http_tool", + description="A test StreamableHttp tool", + inputSchema={ + "type": "object", + "properties": {"test_param": {"type": "string"}}, + "required": ["test_param"], + }, + ) + + @pytest.fixture def mock_sse_session() -> AsyncMock: session = AsyncMock(spec=ClientSession) @@ -53,6 +87,15 @@ def mock_sse_session() -> AsyncMock: return session +@pytest.fixture +def mock_streamable_http_session() -> AsyncMock: + session = AsyncMock(spec=ClientSession) + session.initialize = AsyncMock() + session.call_tool = AsyncMock() + session.list_tools = AsyncMock() + return session + + @pytest.fixture def mock_session() -> AsyncMock: session = AsyncMock(spec=ClientSession) @@ -66,7 +109,13 @@ def mock_session() -> AsyncMock: def mock_tool_response() -> MagicMock: response = MagicMock() response.isError = False - response.content = {"result": "test_output"} + response.content = [ + TextContent( + text="test_output", + type="text", + annotations=Annotations(audience=["user", "assistant"], priority=0.7), + ), + ] return response @@ -75,6 +124,14 @@ def cancellation_token() -> CancellationToken: return CancellationToken() +@pytest.fixture +def mock_error_tool_response() -> MagicMock: + response = MagicMock() + response.isError = True + response.content = [TextContent(text="error output", type="text")] + return response + + def test_adapter_config_serialization(sample_tool: Tool, sample_server_params: StdioServerParams) -> None: """Test that adapter can be saved to and loaded from config.""" original_adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool) @@ -125,7 +182,7 @@ async def test_mcp_tool_execution( with caplog.at_level(logging.INFO): adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool) result = await adapter.run_json( - args=create_model(sample_tool.inputSchema)(**{"test_param": "test"}).model_dump(), + args=schema_to_pydantic_model(sample_tool.inputSchema)(**{"test_param": "test"}).model_dump(), cancellation_token=cancellation_token, ) @@ -177,6 +234,94 @@ async def test_adapter_from_server_params( ) +@pytest.mark.asyncio +async def test_adapter_from_server_params_with_return_value_as_string( + sample_tool: Tool, + sample_server_params: StdioServerParams, + mock_session: AsyncMock, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that adapter can be created from server parameters.""" + mock_context = AsyncMock() + mock_context.__aenter__.return_value = mock_session + monkeypatch.setattr( + "autogen_ext.tools.mcp._base.create_mcp_server_session", + lambda *args, **kwargs: mock_context, # type: ignore + ) + mock_session.list_tools.return_value.tools = [sample_tool] + + adapter = await StdioMcpToolAdapter.from_server_params(sample_server_params, "test_tool") + + assert ( + adapter.return_value_as_string( + [ + TextContent( + text="this is a sample text", + type="text", + annotations=Annotations(audience=["user", "assistant"], priority=0.7), + ), + ImageContent( + data="this is a sample base64 encoded image", + mimeType="image/png", + type="image", + annotations=None, + ), + EmbeddedResource( + type="resource", + resource=TextResourceContents( + text="this is a sample text", + uri=AnyUrl(url="http://example.com/test"), + ), + annotations=Annotations(audience=["user"], priority=0.3), + ), + ] + ) + == '[{"type": "text", "text": "this is a sample text", "annotations": {"audience": ["user", "assistant"], "priority": 0.7}}, {"type": "image", "data": "this is a sample base64 encoded image", "mimeType": "image/png", "annotations": null}, {"type": "resource", "resource": {"uri": "http://example.com/test", "mimeType": null, "text": "this is a sample text"}, "annotations": {"audience": ["user"], "priority": 0.3}}]' + ) + + +@pytest.mark.asyncio +async def test_adapter_from_factory( + sample_tool: Tool, + sample_server_params: StdioServerParams, + mock_session: AsyncMock, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that factory function returns a list of tools.""" + mock_context = AsyncMock() + mock_context.__aenter__.return_value = mock_session + monkeypatch.setattr( + "autogen_ext.tools.mcp._factory.create_mcp_server_session", + lambda *args, **kwargs: mock_context, # type: ignore + ) + mock_session.list_tools.return_value.tools = [sample_tool] + tools = await mcp_server_tools(server_params=sample_server_params) + assert tools is not None + assert len(tools) > 0 + assert isinstance(tools[0], StdioMcpToolAdapter) + + +@pytest.mark.asyncio +async def test_adapter_from_factory_existing_session( + sample_tool: Tool, + sample_server_params: StdioServerParams, + mock_session: AsyncMock, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that factory function returns a list of tools with an existing session.""" + mock_context = AsyncMock() + mock_context.__aenter__.return_value = mock_session + monkeypatch.setattr( + "autogen_ext.tools.mcp._factory.create_mcp_server_session", + lambda *args, **kwargs: mock_context, # type: ignore + ) + mock_session.list_tools.return_value.tools = [sample_tool] + tools = await mcp_server_tools(server_params=sample_server_params, session=mock_session) + assert tools is not None + assert len(tools) > 0 + assert isinstance(tools[0], StdioMcpToolAdapter) + + @pytest.mark.asyncio async def test_sse_adapter_config_serialization(sample_sse_tool: Tool) -> None: """Test that SSE adapter can be saved to and loaded from config.""" @@ -219,7 +364,16 @@ async def test_sse_tool_execution( mock_context = AsyncMock() mock_context.__aenter__.return_value = mock_sse_session - mock_sse_session.call_tool.return_value = MagicMock(isError=False, content={"result": "test_output"}) + mock_sse_session.call_tool.return_value = MagicMock( + isError=False, + content=[ + TextContent( + text="test_output", + type="text", + annotations=Annotations(audience=["user", "assistant"], priority=0.7), + ), + ], + ) monkeypatch.setattr( "autogen_ext.tools.mcp._base.create_mcp_server_session", @@ -229,7 +383,7 @@ async def test_sse_tool_execution( with caplog.at_level(logging.INFO): adapter = SseMcpToolAdapter(server_params=params, tool=sample_sse_tool) result = await adapter.run_json( - args=create_model(sample_sse_tool.inputSchema)(**{"test_param": "test"}).model_dump(), + args=schema_to_pydantic_model(sample_sse_tool.inputSchema)(**{"test_param": "test"}).model_dump(), cancellation_token=CancellationToken(), ) @@ -280,3 +434,497 @@ async def test_sse_adapter_from_server_params( params_schema["properties"]["test_param"]["type"] == sample_sse_tool.inputSchema["properties"]["test_param"]["type"] ) + + +@pytest.mark.asyncio +async def test_streamable_http_adapter_config_serialization(sample_streamable_http_tool: Tool) -> None: + """Test that StreamableHttp adapter can be saved to and loaded from config.""" + params = StreamableHttpServerParams(url="http://test-url") + original_adapter = StreamableHttpMcpToolAdapter(server_params=params, tool=sample_streamable_http_tool) + config = original_adapter.dump_component() + loaded_adapter = StreamableHttpMcpToolAdapter.load_component(config) + + # Test that the loaded adapter has the same properties + assert loaded_adapter.name == "test_streamable_http_tool" + assert loaded_adapter.description == "A test StreamableHttp tool" + + # Verify schema structure + schema = loaded_adapter.schema + assert "parameters" in schema, "Schema must have parameters" + params_schema = schema["parameters"] + assert isinstance(params_schema, dict), "Parameters must be a dict" + assert "type" in params_schema, "Parameters must have type" + assert "required" in params_schema, "Parameters must have required fields" + assert "properties" in params_schema, "Parameters must have properties" + + # Compare schema content + assert params_schema["type"] == sample_streamable_http_tool.inputSchema["type"] + assert params_schema["required"] == sample_streamable_http_tool.inputSchema["required"] + assert ( + params_schema["properties"]["test_param"]["type"] + == sample_streamable_http_tool.inputSchema["properties"]["test_param"]["type"] + ) + + +@pytest.mark.asyncio +async def test_streamable_http_tool_execution( + sample_streamable_http_tool: Tool, + mock_streamable_http_session: AsyncMock, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test that StreamableHttp adapter properly executes tools through ClientSession.""" + params = StreamableHttpServerParams(url="http://test-url") + mock_context = AsyncMock() + mock_context.__aenter__.return_value = mock_streamable_http_session + + mock_streamable_http_session.call_tool.return_value = MagicMock( + isError=False, + content=[ + TextContent( + text="test_output", + type="text", + annotations=Annotations(audience=["user", "assistant"], priority=0.7), + ), + ], + ) + + monkeypatch.setattr( + "autogen_ext.tools.mcp._base.create_mcp_server_session", + lambda *args, **kwargs: mock_context, # type: ignore + ) + + with caplog.at_level(logging.INFO): + adapter = StreamableHttpMcpToolAdapter(server_params=params, tool=sample_streamable_http_tool) + result = await adapter.run_json( + args=schema_to_pydantic_model(sample_streamable_http_tool.inputSchema)( + **{"test_param": "test"} + ).model_dump(), + cancellation_token=CancellationToken(), + ) + + assert result == mock_streamable_http_session.call_tool.return_value.content + mock_streamable_http_session.initialize.assert_called_once() + mock_streamable_http_session.call_tool.assert_called_once() + + # Check log. + assert "test_output" in caplog.text + + +@pytest.mark.asyncio +async def test_streamable_http_adapter_from_server_params( + sample_streamable_http_tool: Tool, + mock_streamable_http_session: AsyncMock, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that StreamableHttp adapter can be created from server parameters.""" + params = StreamableHttpServerParams(url="http://test-url") + mock_context = AsyncMock() + mock_context.__aenter__.return_value = mock_streamable_http_session + monkeypatch.setattr( + "autogen_ext.tools.mcp._base.create_mcp_server_session", + lambda *args, **kwargs: mock_context, # type: ignore + ) + + mock_streamable_http_session.list_tools.return_value.tools = [sample_streamable_http_tool] + + adapter = await StreamableHttpMcpToolAdapter.from_server_params(params, "test_streamable_http_tool") + + assert isinstance(adapter, StreamableHttpMcpToolAdapter) + assert adapter.name == "test_streamable_http_tool" + assert adapter.description == "A test StreamableHttp tool" + + # Verify schema structure + schema = adapter.schema + assert "parameters" in schema, "Schema must have parameters" + params_schema = schema["parameters"] + assert isinstance(params_schema, dict), "Parameters must be a dict" + assert "type" in params_schema, "Parameters must have type" + assert "required" in params_schema, "Parameters must have required fields" + assert "properties" in params_schema, "Parameters must have properties" + + # Compare schema content + assert params_schema["type"] == sample_streamable_http_tool.inputSchema["type"] + assert params_schema["required"] == sample_streamable_http_tool.inputSchema["required"] + assert ( + params_schema["properties"]["test_param"]["type"] + == sample_streamable_http_tool.inputSchema["properties"]["test_param"]["type"] + ) + + +@pytest.mark.asyncio +async def test_mcp_server_fetch() -> None: + params = StdioServerParams( + command="uvx", + args=["mcp-server-fetch"], + read_timeout_seconds=60, + ) + tools = await mcp_server_tools(server_params=params) + assert tools is not None + assert tools[0].name == "fetch" + result = await tools[0].run_json({"url": "https://github.com/"}, CancellationToken()) + assert result is not None + + +@pytest.mark.asyncio +async def test_mcp_server_filesystem() -> None: + params = StdioServerParams( + command="npx", + args=[ + "-y", + "@modelcontextprotocol/server-filesystem", + ".", + ], + read_timeout_seconds=60, + ) + tools = await mcp_server_tools(server_params=params) + assert tools is not None + tools = [tool for tool in tools if tool.name == "read_file"] + assert len(tools) == 1 + tool = tools[0] + result = await tool.run_json({"path": "README.md"}, CancellationToken()) + assert result is not None + + +@pytest.mark.asyncio +async def test_mcp_server_git() -> None: + params = StdioServerParams( + command="uvx", + args=["mcp-server-git"], + read_timeout_seconds=60, + ) + tools = await mcp_server_tools(server_params=params) + assert tools is not None + tools = [tool for tool in tools if tool.name == "git_log"] + assert len(tools) == 1 + tool = tools[0] + repo_path = os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "..") + result = await tool.run_json({"repo_path": repo_path}, CancellationToken()) + assert result is not None + + +@pytest.mark.asyncio +async def test_mcp_server_git_existing_session() -> None: + params = StdioServerParams( + command="uvx", + args=["mcp-server-git"], + read_timeout_seconds=60, + ) + async with create_mcp_server_session(params) as session: + await session.initialize() + tools = await mcp_server_tools(server_params=params, session=session) + assert tools is not None + git_log = [tool for tool in tools if tool.name == "git_log"][0] + repo_path = os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "..") + result = await git_log.run_json({"repo_path": repo_path}, CancellationToken()) + assert result is not None + + git_status = [tool for tool in tools if tool.name == "git_status"][0] + result = await git_status.run_json({"repo_path": repo_path}, CancellationToken()) + assert result is not None + + +@pytest.mark.asyncio +async def test_mcp_server_github() -> None: + # Check if GITHUB_TOKEN is set. + if "GITHUB_TOKEN" not in os.environ: + pytest.skip("GITHUB_TOKEN environment variable is not set. Skipping test.") + params = StdioServerParams( + command="npx", + args=[ + "-y", + "@modelcontextprotocol/server-github", + ], + env={"GITHUB_PERSONAL_ACCESS_TOKEN": os.environ["GITHUB_TOKEN"]}, + read_timeout_seconds=60, + ) + tools = await mcp_server_tools(server_params=params) + assert tools is not None + tools = [tool for tool in tools if tool.name == "get_file_contents"] + assert len(tools) == 1 + tool = tools[0] + result = await tool.run_json( + {"owner": "microsoft", "repo": "autogen", "path": "python", "branch": "main"}, + CancellationToken(), + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_mcp_workbench_start_stop() -> None: + params = StdioServerParams( + command="uvx", + args=["mcp-server-fetch"], + read_timeout_seconds=60, + ) + + workbench = McpWorkbench(params) + assert workbench is not None + assert workbench.server_params == params + await workbench.start() + assert workbench._actor is not None # type: ignore[reportPrivateUsage] + await workbench.stop() + assert workbench._actor is None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_mcp_workbench_server_fetch() -> None: + params = StdioServerParams( + command="uvx", + args=["mcp-server-fetch"], + read_timeout_seconds=60, + ) + + workbench = McpWorkbench(server_params=params) + await workbench.start() + + tools = await workbench.list_tools() + assert tools is not None + assert tools[0]["name"] == "fetch" + + result = await workbench.call_tool(tools[0]["name"], {"url": "https://github.com/"}, CancellationToken()) + assert result is not None + + await workbench.stop() + + +@pytest.mark.asyncio +async def test_mcp_workbench_server_filesystem() -> None: + params = StdioServerParams( + command="npx", + args=[ + "-y", + "@modelcontextprotocol/server-filesystem", + ".", + ], + read_timeout_seconds=60, + ) + + workbench = McpWorkbench(server_params=params) + await workbench.start() + + tools = await workbench.list_tools() + assert tools is not None + tools = [tool for tool in tools if tool["name"] == "read_file"] + assert len(tools) == 1 + tool = tools[0] + result = await workbench.call_tool(tool["name"], {"path": "README.md"}, CancellationToken()) + assert result is not None + + await workbench.stop() + + # Serialize the workbench. + config = workbench.dump_component() + + # Deserialize the workbench. + async with Workbench.load_component(config) as new_workbench: + tools = await new_workbench.list_tools() + assert tools is not None + tools = [tool for tool in tools if tool["name"] == "read_file"] + assert len(tools) == 1 + tool = tools[0] + result = await new_workbench.call_tool(tool["name"], {"path": "README.md"}, CancellationToken()) + assert result is not None + + +@pytest.mark.asyncio +async def test_lazy_init_and_finalize_cleanup() -> None: + params = StdioServerParams( + command="npx", + args=[ + "-y", + "@modelcontextprotocol/server-filesystem", + ".", + ], + read_timeout_seconds=60, + ) + workbench = McpWorkbench(server_params=params) + + # Before any call, actor should not be initialized + assert workbench._actor is None # type: ignore[reportPrivateUsage] + + # Trigger list_tools → lazy start + await workbench.list_tools() + assert workbench._actor is not None # type: ignore[reportPrivateUsage] + assert workbench._actor._active is True # type: ignore[reportPrivateUsage] + + actor = workbench._actor # type: ignore[reportPrivateUsage] + del workbench + await asyncio.sleep(0.1) + assert actor._active is False + + +@pytest.mark.asyncio +async def test_del_to_new_event_loop_when_get_event_loop_fails() -> None: + params = StdioServerParams( + command="npx", + args=[ + "-y", + "@modelcontextprotocol/server-filesystem", + ".", + ], + read_timeout_seconds=60, + ) + workbench = McpWorkbench(server_params=params) + + await workbench.list_tools() + assert workbench._actor is not None # type: ignore[reportPrivateUsage] + assert workbench._actor._active is True # type: ignore[reportPrivateUsage] + + actor = workbench._actor # type: ignore[reportPrivateUsage] + + def cleanup() -> None: + nonlocal workbench + del workbench + + t = threading.Thread(target=cleanup) + t.start() + t.join() + + await asyncio.sleep(0.1) + assert actor._active is False # type: ignore[reportPrivateUsage] + + +def test_del_raises_when_loop_closed() -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + params = StdioServerParams(command="echo", args=["ok"]) + workbench = McpWorkbench(server_params=params) + + workbench._actor_loop = loop # type: ignore[reportPrivateUsage] + workbench._actor = cast(McpSessionActor, object()) # type: ignore[reportPrivateUsage] + + loop.close() + + with pytest.warns(RuntimeWarning, match="loop is closed or not running"): + del workbench + + +def test_mcp_tool_adapter_normalize_payload(sample_tool: Tool, sample_server_params: StdioServerParams) -> None: + """Test the _normalize_payload_to_content_list method of McpToolAdapter.""" + adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool) + + # Case 1: Payload is already a list of valid content items + valid_content_list: list[TextContent | ImageContent | EmbeddedResource] = [ + TextContent(text="hello", type="text"), + ImageContent(data="base64data", mimeType="image/png", type="image"), + EmbeddedResource( + type="resource", + resource=TextResourceContents(text="embedded text", uri=AnyUrl(url="http://example.com/resource")), + ), + ] + assert adapter._normalize_payload_to_content_list(valid_content_list) == valid_content_list # type: ignore[reportPrivateUsage] + + # Case 2: Payload is a single TextContent + single_text_content = TextContent(text="single text", type="text") + assert adapter._normalize_payload_to_content_list(single_text_content) == [single_text_content] # type: ignore[reportPrivateUsage, arg-type] + + # Case 3: Payload is a single ImageContent + single_image_content = ImageContent(data="imagedata", mimeType="image/jpeg", type="image") + assert adapter._normalize_payload_to_content_list(single_image_content) == [single_image_content] # type: ignore[reportPrivateUsage, arg-type] + + # Case 4: Payload is a single EmbeddedResource + single_embedded_resource = EmbeddedResource( + type="resource", + resource=TextResourceContents(text="other embedded", uri=AnyUrl(url="http://example.com/other")), + ) + assert adapter._normalize_payload_to_content_list(single_embedded_resource) == [single_embedded_resource] # type: ignore[reportPrivateUsage, arg-type] + + # Case 5: Payload is a string + string_payload = "This is a string payload." + expected_from_string = [TextContent(text=string_payload, type="text")] + assert adapter._normalize_payload_to_content_list(string_payload) == expected_from_string # type: ignore[reportPrivateUsage, arg-type] + + # Case 6: Payload is an integer + int_payload = 12345 + expected_from_int = [TextContent(text=str(int_payload), type="text")] + assert adapter._normalize_payload_to_content_list(int_payload) == expected_from_int # type: ignore[reportPrivateUsage, arg-type] + + # Case 7: Payload is a dictionary + dict_payload = {"key": "value", "number": 42} + expected_from_dict = [TextContent(text=str(dict_payload), type="text")] + assert adapter._normalize_payload_to_content_list(dict_payload) == expected_from_dict # type: ignore[reportPrivateUsage, arg-type] + + # Case 8: Payload is an empty list (should still be a list of valid items, so returns as is) + empty_list_payload: list[TextContent | ImageContent | EmbeddedResource] = [] + assert adapter._normalize_payload_to_content_list(empty_list_payload) == empty_list_payload # type: ignore[reportPrivateUsage] + + # Case 9: Payload is None (should be stringified) + none_payload = None + expected_from_none = [TextContent(text=str(none_payload), type="text")] + assert adapter._normalize_payload_to_content_list(none_payload) == expected_from_none # type: ignore[reportPrivateUsage, arg-type] + + +@pytest.mark.asyncio +async def test_mcp_tool_adapter_run_error( + sample_tool: Tool, + sample_server_params: StdioServerParams, + mock_session: AsyncMock, + mock_error_tool_response: MagicMock, + cancellation_token: CancellationToken, +) -> None: + """Test McpToolAdapter._run when tool returns an error.""" + adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool, session=mock_session) + mock_session.call_tool.return_value = mock_error_tool_response + + args = {"test_param": "test_value"} + with pytest.raises(Exception) as excinfo: + await adapter._run(args=args, cancellation_token=cancellation_token, session=mock_session) # type: ignore[reportPrivateUsage] + + mock_session.call_tool.assert_called_once_with(name=sample_tool.name, arguments=args) + assert adapter.return_value_as_string([TextContent(text="error output", type="text")]) in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_mcp_tool_adapter_run_cancelled_before_call( + sample_tool: Tool, + sample_server_params: StdioServerParams, + mock_session: AsyncMock, + cancellation_token: CancellationToken, +) -> None: + """Test McpToolAdapter._run when operation is cancelled before tool call.""" + adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool, session=mock_session) + cancellation_token.cancel() # Cancel before the call + + args = {"test_param": "test_value"} + with pytest.raises(asyncio.CancelledError): + await adapter._run(args=args, cancellation_token=cancellation_token, session=mock_session) # type: ignore[reportPrivateUsage] + + mock_session.call_tool.assert_not_called() + + +@pytest.mark.asyncio +async def test_mcp_tool_adapter_run_cancelled_during_call( + sample_tool: Tool, + sample_server_params: StdioServerParams, + mock_session: AsyncMock, + cancellation_token: CancellationToken, +) -> None: + """Test McpToolAdapter._run when operation is cancelled during tool call.""" + adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool, session=mock_session) + mock_session.call_tool.side_effect = asyncio.CancelledError("Tool call cancelled") + + args = {"test_param": "test_value"} + with pytest.raises(asyncio.CancelledError): + await adapter._run(args=args, cancellation_token=cancellation_token, session=mock_session) # type: ignore[reportPrivateUsage] + + mock_session.call_tool.assert_called_once_with(name=sample_tool.name, arguments=args) + + +def test_return_value_as_string_with_resource_link(sample_tool: Tool, sample_server_params: StdioServerParams) -> None: + """Test return_value_as_string handles ResourceLink objects correctly.""" + adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool) + + # Test ResourceLink with meta field + resource_link = ResourceLink( + name="test_link", + type="resource_link", + uri=AnyUrl(url="http://example.com"), + ) + + result = adapter.return_value_as_string([resource_link]) + # Verify the JSON serialization contains expected fields + assert '"type": "resource_link"' in result + assert '"name": "test_link"' in result + assert '"uri": "http://example.com/"' in result # AnyUrl normalizes with trailing slash diff --git a/python/packages/autogen-ext/tests/tools/test_mcp_workbench_features.py b/python/packages/autogen-ext/tests/tools/test_mcp_workbench_features.py new file mode 100644 index 000000000000..6f08633da11a --- /dev/null +++ b/python/packages/autogen-ext/tests/tools/test_mcp_workbench_features.py @@ -0,0 +1,551 @@ +import asyncio +from unittest.mock import AsyncMock + +import pytest +from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams +from mcp.types import ( + GetPromptResult, + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, + Prompt, + PromptArgument, + PromptMessage, + ReadResourceResult, + Resource, + ResourceTemplate, + TextContent, + TextResourceContents, +) +from pydantic import AnyUrl + + +@pytest.fixture +def sample_server_params() -> StdioServerParams: + """Sample server parameters for testing.""" + return StdioServerParams(command="echo", args=["test"]) + + +@pytest.fixture +def mock_mcp_actor() -> AsyncMock: + """Mock MCP session actor.""" + actor = AsyncMock() + return actor + + +@pytest.fixture +def sample_prompts() -> list[Prompt]: + """Create sample MCP prompts for testing.""" + return [ + Prompt( + name="code_review", + description="Reviews code for best practices", + arguments=[ + PromptArgument( + name="language", + description="Programming language", + required=True, + ) + ], + ), + Prompt( + name="documentation", + description="Generates documentation", + arguments=[ + PromptArgument( + name="format", + description="Output format", + required=False, + ) + ], + ), + ] + + +@pytest.fixture +def sample_resources() -> list[Resource]: + """Create sample MCP resources for testing.""" + return [ + Resource( + uri=AnyUrl("file:///test/document.txt"), + name="Test Document", + description="A sample document for testing", + mimeType="text/plain", + ), + Resource( + uri=AnyUrl("https://api.example.com/data"), + name="API Data", + description="External API data source", + mimeType="application/json", + ), + ] + + +@pytest.fixture +def sample_resource_templates() -> list[ResourceTemplate]: + """Create sample MCP resource templates for testing.""" + return [ + ResourceTemplate( + uriTemplate="file:///logs/{date}.log", + name="Daily Logs", + description="Daily log files by date", + mimeType="text/plain", + ), + ResourceTemplate( + uriTemplate="https://api.example.com/users/{userId}", + name="User Profile", + description="User profile by ID", + mimeType="application/json", + ), + ] + + +@pytest.mark.asyncio +async def test_list_prompts( + sample_prompts: list[Prompt], mock_mcp_actor: AsyncMock, sample_server_params: StdioServerParams +) -> None: + """Test listing prompts from MCP server.""" + # Create workbench + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = mock_mcp_actor # type: ignore[reportPrivateUsage] + + # Mock list_prompts response + list_prompts_result = ListPromptsResult(prompts=sample_prompts) + future_result: asyncio.Future[ListPromptsResult] = asyncio.Future() + future_result.set_result(list_prompts_result) + mock_mcp_actor.call.return_value = future_result + + try: + # List prompts + result = await workbench.list_prompts() + + # Verify result + assert isinstance(result, ListPromptsResult) + assert len(result.prompts) == 2 + assert result.prompts[0].name == "code_review" + assert result.prompts[0].description == "Reviews code for best practices" + assert result.prompts[1].name == "documentation" + assert result.prompts[1].description == "Generates documentation" + + # Verify actor was called correctly + mock_mcp_actor.call.assert_called_with("list_prompts", None) + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_get_prompt_without_arguments(mock_mcp_actor: AsyncMock, sample_server_params: StdioServerParams) -> None: + """Test getting a prompt without arguments.""" + # Create workbench + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = mock_mcp_actor # type: ignore[reportPrivateUsage] + + # Mock get_prompt response + get_prompt_result = GetPromptResult( + description="Code review prompt", + messages=[ + PromptMessage( + role="user", + content=TextContent( + type="text", + text="Please review this code for best practices and suggest improvements.", + ), + ) + ], + ) + future_result: asyncio.Future[GetPromptResult] = asyncio.Future() + future_result.set_result(get_prompt_result) + mock_mcp_actor.call.return_value = future_result + + try: + # Get prompt + result = await workbench.get_prompt("code_review") + + # Verify result + assert isinstance(result, GetPromptResult) + assert result.description == "Code review prompt" + assert len(result.messages) == 1 + assert result.messages[0].role == "user" + + # Verify actor was called correctly + mock_mcp_actor.call.assert_called_with("get_prompt", {"name": "code_review", "kargs": {"arguments": None}}) + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_get_prompt_with_arguments(mock_mcp_actor: AsyncMock, sample_server_params: StdioServerParams) -> None: + """Test getting a prompt with arguments.""" + # Create workbench + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = mock_mcp_actor # type: ignore[reportPrivateUsage] + + # Mock get_prompt response + get_prompt_result = GetPromptResult( + description="Python code review prompt", + messages=[ + PromptMessage( + role="user", + content=TextContent( + type="text", + text="Please review this Python code for best practices and suggest improvements.", + ), + ) + ], + ) + future_result: asyncio.Future[GetPromptResult] = asyncio.Future() + future_result.set_result(get_prompt_result) + mock_mcp_actor.call.return_value = future_result + + try: + # Get prompt with arguments + arguments = {"language": "python", "style": "pep8"} + result = await workbench.get_prompt("code_review", arguments) + + # Verify result + assert isinstance(result, GetPromptResult) + assert result.description == "Python code review prompt" + assert len(result.messages) == 1 + + # Verify actor was called correctly + mock_mcp_actor.call.assert_called_with("get_prompt", {"name": "code_review", "kargs": {"arguments": arguments}}) + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_list_resources( + sample_resources: list[Resource], mock_mcp_actor: AsyncMock, sample_server_params: StdioServerParams +) -> None: + """Test listing resources from MCP server.""" + # Create workbench + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = mock_mcp_actor # type: ignore[reportPrivateUsage] + + # Mock list_resources response + list_resources_result = ListResourcesResult(resources=sample_resources) + future_result: asyncio.Future[ListResourcesResult] = asyncio.Future() + future_result.set_result(list_resources_result) + mock_mcp_actor.call.return_value = future_result + + try: + # List resources + result = await workbench.list_resources() + + # Verify result + assert isinstance(result, ListResourcesResult) + assert len(result.resources) == 2 + assert str(result.resources[0].uri) == "file:///test/document.txt" + assert result.resources[0].name == "Test Document" + assert result.resources[0].mimeType == "text/plain" + assert str(result.resources[1].uri) == "https://api.example.com/data" + assert result.resources[1].name == "API Data" + assert result.resources[1].mimeType == "application/json" + + # Verify actor was called correctly + mock_mcp_actor.call.assert_called_with("list_resources", None) + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_read_resource(mock_mcp_actor: AsyncMock, sample_server_params: StdioServerParams) -> None: + """Test reading a resource from MCP server.""" + # Create workbench + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = mock_mcp_actor # type: ignore[reportPrivateUsage] + + # Mock read_resource response + read_resource_result = ReadResourceResult( + contents=[ + TextResourceContents( + uri=AnyUrl("file:///test/document.txt"), + mimeType="text/plain", + text="This is the content of the test document.", + ) + ] + ) + future_result: asyncio.Future[ReadResourceResult] = asyncio.Future() + future_result.set_result(read_resource_result) + mock_mcp_actor.call.return_value = future_result + + try: + # Read resource + uri = "file:///test/document.txt" + result = await workbench.read_resource(uri) + + # Verify result + assert isinstance(result, ReadResourceResult) + assert len(result.contents) == 1 + content = result.contents[0] + assert isinstance(content, TextResourceContents) + assert content.uri == AnyUrl(uri) + assert content.mimeType == "text/plain" + assert content.text == "This is the content of the test document." + + # Verify actor was called correctly + mock_mcp_actor.call.assert_called_with("read_resource", {"name": None, "kargs": {"uri": uri}}) + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_list_resource_templates( + sample_resource_templates: list[ResourceTemplate], + mock_mcp_actor: AsyncMock, + sample_server_params: StdioServerParams, +) -> None: + """Test listing resource templates from MCP server.""" + # Create workbench + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = mock_mcp_actor # type: ignore[reportPrivateUsage] + + # Mock list_resource_templates response + list_templates_result = ListResourceTemplatesResult(resourceTemplates=sample_resource_templates) + future_result: asyncio.Future[ListResourceTemplatesResult] = asyncio.Future() + future_result.set_result(list_templates_result) + mock_mcp_actor.call.return_value = future_result + + try: + # List resource templates + result = await workbench.list_resource_templates() + + # Verify result + assert isinstance(result, ListResourceTemplatesResult) + assert len(result.resourceTemplates) == 2 + assert result.resourceTemplates[0].uriTemplate == "file:///logs/{date}.log" + assert result.resourceTemplates[0].name == "Daily Logs" + assert result.resourceTemplates[0].mimeType == "text/plain" + assert result.resourceTemplates[1].uriTemplate == "https://api.example.com/users/{userId}" + assert result.resourceTemplates[1].name == "User Profile" + assert result.resourceTemplates[1].mimeType == "application/json" + + # Verify actor was called correctly + mock_mcp_actor.call.assert_called_with("list_resource_templates", None) + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_workbench_auto_start_on_list_prompts( + sample_prompts: list[Prompt], sample_server_params: StdioServerParams +) -> None: + """Test that workbench automatically starts when list_prompts is called without explicit start.""" + # Create workbench without starting it + workbench = McpWorkbench(server_params=sample_server_params) + + # Mock the start method to avoid actual server connection + original_start = workbench.start + start_called = False + + async def mock_start() -> None: + nonlocal start_called + start_called = True + # Set up a mock actor + mock_actor = AsyncMock() + list_prompts_result = ListPromptsResult(prompts=sample_prompts) + future_result: asyncio.Future[ListPromptsResult] = asyncio.Future() + future_result.set_result(list_prompts_result) + mock_actor.call.return_value = future_result + workbench._actor = mock_actor # type: ignore[reportPrivateUsage] + + workbench.start = mock_start # type: ignore[method-assign] + + try: + # Call list_prompts without explicitly starting + result = await workbench.list_prompts() + + # Verify that start was called + assert start_called + assert isinstance(result, ListPromptsResult) + assert len(result.prompts) == 2 + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + workbench.start = original_start # type: ignore[method-assign] + + +@pytest.mark.asyncio +async def test_workbench_auto_start_on_get_prompt(sample_server_params: StdioServerParams) -> None: + """Test that workbench automatically starts when get_prompt is called without explicit start.""" + # Create workbench without starting it + workbench = McpWorkbench(server_params=sample_server_params) + + # Mock the start method to avoid actual server connection + start_called = False + + async def mock_start() -> None: + nonlocal start_called + start_called = True + # Set up a mock actor + mock_actor = AsyncMock() + get_prompt_result = GetPromptResult( + description="Test prompt", + messages=[PromptMessage(role="user", content=TextContent(type="text", text="Test message"))], + ) + future_result: asyncio.Future[GetPromptResult] = asyncio.Future() + future_result.set_result(get_prompt_result) + mock_actor.call.return_value = future_result + workbench._actor = mock_actor # type: ignore[reportPrivateUsage] + + workbench.start = mock_start # type: ignore[method-assign] + + try: + # Call get_prompt without explicitly starting + result = await workbench.get_prompt("test_prompt") + + # Verify that start was called + assert start_called + assert isinstance(result, GetPromptResult) + assert result.description == "Test prompt" + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_workbench_auto_start_on_list_resources( + sample_resources: list[Resource], sample_server_params: StdioServerParams +) -> None: + """Test that workbench automatically starts when list_resources is called without explicit start.""" + # Create workbench without starting it + workbench = McpWorkbench(server_params=sample_server_params) + + # Mock the start method to avoid actual server connection + start_called = False + + async def mock_start() -> None: + nonlocal start_called + start_called = True + # Set up a mock actor + mock_actor = AsyncMock() + list_resources_result = ListResourcesResult(resources=sample_resources) + future_result: asyncio.Future[ListResourcesResult] = asyncio.Future() + future_result.set_result(list_resources_result) + mock_actor.call.return_value = future_result + workbench._actor = mock_actor # type: ignore[reportPrivateUsage] + + workbench.start = mock_start # type: ignore[method-assign] + + try: + # Call list_resources without explicitly starting + result = await workbench.list_resources() + + # Verify that start was called + assert start_called + assert isinstance(result, ListResourcesResult) + assert len(result.resources) == 2 + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_workbench_auto_start_on_read_resource(sample_server_params: StdioServerParams) -> None: + """Test that workbench automatically starts when read_resource is called without explicit start.""" + # Create workbench without starting it + workbench = McpWorkbench(server_params=sample_server_params) + + # Mock the start method to avoid actual server connection + start_called = False + + async def mock_start() -> None: + nonlocal start_called + start_called = True + # Set up a mock actor + mock_actor = AsyncMock() + read_resource_result = ReadResourceResult( + contents=[TextResourceContents(uri=AnyUrl("file:///test.txt"), mimeType="text/plain", text="Test content")] + ) + future_result: asyncio.Future[ReadResourceResult] = asyncio.Future() + future_result.set_result(read_resource_result) + mock_actor.call.return_value = future_result + workbench._actor = mock_actor # type: ignore[reportPrivateUsage] + + workbench.start = mock_start # type: ignore[method-assign] + + try: + # Call read_resource without explicitly starting + result = await workbench.read_resource("file:///test.txt") + + # Verify that start was called + assert start_called + assert isinstance(result, ReadResourceResult) + assert len(result.contents) == 1 + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_workbench_auto_start_on_list_resource_templates( + sample_resource_templates: list[ResourceTemplate], sample_server_params: StdioServerParams +) -> None: + """Test that workbench automatically starts when list_resource_templates is called without explicit start.""" + # Create workbench without starting it + workbench = McpWorkbench(server_params=sample_server_params) + + # Mock the start method to avoid actual server connection + start_called = False + + async def mock_start() -> None: + nonlocal start_called + start_called = True + # Set up a mock actor + mock_actor = AsyncMock() + list_templates_result = ListResourceTemplatesResult(resourceTemplates=sample_resource_templates) + future_result: asyncio.Future[ListResourceTemplatesResult] = asyncio.Future() + future_result.set_result(list_templates_result) + mock_actor.call.return_value = future_result + workbench._actor = mock_actor # type: ignore[reportPrivateUsage] + + workbench.start = mock_start # type: ignore[method-assign] + + try: + # Call list_resource_templates without explicitly starting + result = await workbench.list_resource_templates() + + # Verify that start was called + assert start_called + assert isinstance(result, ListResourceTemplatesResult) + assert len(result.resourceTemplates) == 2 + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_workbench_methods_raise_error_if_actor_fails_to_initialize( + sample_server_params: StdioServerParams, +) -> None: + """Test that methods raise RuntimeError if actor fails to initialize.""" + # Create workbench + workbench = McpWorkbench(server_params=sample_server_params) + + # Mock start to set actor to None (simulating initialization failure) + async def mock_start() -> None: + workbench._actor = None # type: ignore[reportPrivateUsage] + + workbench.start = mock_start # type: ignore[method-assign] + + # Test that all methods raise RuntimeError when actor is None + with pytest.raises(RuntimeError, match="Actor is not initialized"): + await workbench.list_prompts() + + with pytest.raises(RuntimeError, match="Actor is not initialized"): + await workbench.get_prompt("test") + + with pytest.raises(RuntimeError, match="Actor is not initialized"): + await workbench.list_resources() + + with pytest.raises(RuntimeError, match="Actor is not initialized"): + await workbench.read_resource("file:///test.txt") + + with pytest.raises(RuntimeError, match="Actor is not initialized"): + await workbench.list_resource_templates() diff --git a/python/packages/autogen-ext/tests/tools/test_mcp_workbench_model_client.py b/python/packages/autogen-ext/tests/tools/test_mcp_workbench_model_client.py new file mode 100644 index 000000000000..fe6567bf7503 --- /dev/null +++ b/python/packages/autogen-ext/tests/tools/test_mcp_workbench_model_client.py @@ -0,0 +1,831 @@ +import asyncio +from typing import Any, AsyncGenerator, Sequence +from unittest.mock import AsyncMock, MagicMock + +import pytest +from autogen_core import Component +from autogen_core.models import ( + ChatCompletionClient, + CreateResult, + LLMMessage, + ModelInfo, + RequestUsage, +) +from autogen_core.tools import Tool, ToolSchema +from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams +from mcp.types import CallToolResult, ListToolsResult, TextContent +from mcp.types import Tool as McpTool +from pydantic import BaseModel + + +class MockChatCompletionClientConfig(BaseModel): + model_info: ModelInfo + + +class MockChatCompletionClient(ChatCompletionClient, Component[MockChatCompletionClientConfig]): + """Mock chat completion client for testing.""" + + component_type = "test" + component_config_schema = MockChatCompletionClientConfig + component_provider_override = f"{__module__}.MockChatCompletionClient" + + def __init__(self, model_info: ModelInfo | None = None): + self._model_info = model_info or { + "vision": False, + "function_calling": False, + "json_output": False, + "family": "test-model", + "structured_output": False, + } + + @property + def model_info(self) -> ModelInfo: + return self._model_info + + async def create(self, *args: Any, **kwargs: Any) -> CreateResult: + return CreateResult( + content="Mock response", + finish_reason="stop", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + + def create_stream(self, *args: Any, **kwargs: Any) -> AsyncGenerator[str | CreateResult, None]: + async def mock_stream() -> AsyncGenerator[str | CreateResult, None]: + yield "Mock streaming response" + yield CreateResult( + content="Mock response", + finish_reason="stop", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + + return mock_stream() + + async def close(self) -> None: + pass + + def actual_usage(self) -> RequestUsage: + return RequestUsage(prompt_tokens=0, completion_tokens=0) + + def total_usage(self) -> RequestUsage: + return RequestUsage(prompt_tokens=0, completion_tokens=0) + + def count_tokens(self, messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = []) -> int: + return 10 + + def remaining_tokens(self, messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = []) -> int: + return 1000 + + @property + def capabilities(self) -> Any: # ModelCapabilities is deprecated + return {} + + def _to_config(self) -> BaseModel: + return MockChatCompletionClientConfig(model_info=self.model_info) + + @classmethod + def _from_config(cls, config: MockChatCompletionClientConfig) -> "MockChatCompletionClient": + return cls(config.model_info) + + +@pytest.fixture +def sample_mcp_tools() -> list[McpTool]: + """Create sample MCP tools for testing.""" + return [ + McpTool( + name="test_tool", + description="A test tool", + inputSchema={ + "type": "object", + "properties": {"param": {"type": "string"}}, + "required": ["param"], + }, + ), + ] + + +@pytest.fixture +def mock_mcp_actor() -> AsyncMock: + """Mock MCP session actor.""" + actor = AsyncMock() + return actor + + +@pytest.fixture +def sample_server_params() -> StdioServerParams: + """Sample server parameters for testing.""" + return StdioServerParams(command="echo", args=["test"]) + + +@pytest.fixture +def mock_model_client() -> MockChatCompletionClient: + """Mock model client for testing.""" + model_client = MockChatCompletionClient() + # Set up the mock's create method + mock_create = AsyncMock( + return_value=CreateResult( + content="Mock response", + finish_reason="stop", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ) + model_client.create = mock_create # type: ignore[method-assign] + return model_client + + +@pytest.fixture +def mock_model_client_with_vision() -> MockChatCompletionClient: + """Mock model client with vision support for testing.""" + model_client = MockChatCompletionClient( + model_info={ + "vision": True, + "function_calling": False, + "json_output": False, + "family": "test-vision-model", + "structured_output": False, + } + ) + # Set up the mock's create method + mock_create = AsyncMock( + return_value=CreateResult( + content="Mock response", + finish_reason="stop", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ) + model_client.create = mock_create # type: ignore[method-assign] + return model_client + + +@pytest.mark.asyncio +async def test_mcp_workbench_with_model_client_initialization( + sample_server_params: StdioServerParams, + mock_model_client: MockChatCompletionClient, +) -> None: + """Test McpWorkbench initialization with a model client.""" + workbench = McpWorkbench( + server_params=sample_server_params, + model_client=mock_model_client, + ) + + assert workbench._model_client is mock_model_client # type: ignore[attr-defined] + + # Test that the workbench is properly configured + config = workbench._to_config() # type: ignore[attr-defined] + assert config.model_client is not None + assert config.server_params == sample_server_params + + +@pytest.mark.asyncio +async def test_mcp_workbench_without_model_client_initialization( + sample_server_params: StdioServerParams, +) -> None: + """Test McpWorkbench initialization without a model client.""" + workbench = McpWorkbench(server_params=sample_server_params) + + assert workbench._model_client is None # type: ignore[attr-defined] + + # Test that the workbench is properly configured + config = workbench._to_config() # type: ignore[attr-defined] + assert config.model_client is None + assert config.server_params == sample_server_params + + +@pytest.mark.asyncio +async def test_mcp_workbench_serialization_with_model_client( + sample_server_params: StdioServerParams, + mock_model_client: MockChatCompletionClient, +) -> None: + """Test serialization and deserialization of McpWorkbench with model client.""" + # Create original workbench + original_workbench = McpWorkbench( + server_params=sample_server_params, + model_client=mock_model_client, + ) + + # Serialize to config + config = original_workbench._to_config() # type: ignore[attr-defined] + assert config.model_client is not None + + # Deserialize from config + loaded_workbench = McpWorkbench._from_config(config) # type: ignore[attr-defined] + + # Verify the loaded workbench + assert loaded_workbench.server_params == sample_server_params + assert loaded_workbench._model_client is not None # type: ignore[attr-defined] + assert loaded_workbench._model_client.model_info == mock_model_client.model_info # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_mcp_workbench_serialization_without_model_client( + sample_server_params: StdioServerParams, +) -> None: + """Test serialization and deserialization of McpWorkbench without model client.""" + # Create original workbench + original_workbench = McpWorkbench(server_params=sample_server_params) + + # Serialize to config + config = original_workbench._to_config() # type: ignore[attr-defined] + assert config.model_client is None + + # Deserialize from config + loaded_workbench = McpWorkbench._from_config(config) # type: ignore[attr-defined] + + # Verify the loaded workbench + assert loaded_workbench.server_params == sample_server_params + assert loaded_workbench._model_client is None # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_mcp_workbench_model_client_passed_to_actor( + sample_mcp_tools: list[McpTool], + mock_mcp_actor: AsyncMock, + sample_server_params: StdioServerParams, + mock_model_client: MockChatCompletionClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that model client is properly passed to the MCP actor.""" + # Mock the McpSessionActor creation + mock_actor_class = MagicMock() + mock_actor_instance = mock_mcp_actor + mock_actor_class.return_value = mock_actor_instance + + monkeypatch.setattr("autogen_ext.tools.mcp._workbench.McpSessionActor", mock_actor_class) + + # Create workbench with model client + workbench = McpWorkbench( + server_params=sample_server_params, + model_client=mock_model_client, + ) + + # Mock list_tools response + list_tools_result = ListToolsResult(tools=sample_mcp_tools) + future_result: asyncio.Future[ListToolsResult] = asyncio.Future() + future_result.set_result(list_tools_result) + mock_mcp_actor.call.return_value = future_result + + try: + # Start the workbench - this should create the actor with the model client + await workbench.start() + + # Verify that the actor was created with the model client + mock_actor_class.assert_called_once_with(sample_server_params, model_client=mock_model_client) + + # Verify that initialize was called on the actor + mock_mcp_actor.initialize.assert_called_once() + + finally: + if workbench._actor: # type: ignore[attr-defined] + workbench._actor = None # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_mcp_workbench_start_without_explicit_model_client( + sample_mcp_tools: list[McpTool], + mock_mcp_actor: AsyncMock, + sample_server_params: StdioServerParams, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that actor is created correctly when no model client is provided.""" + # Mock the McpSessionActor creation + mock_actor_class = MagicMock() + mock_actor_instance = mock_mcp_actor + mock_actor_class.return_value = mock_actor_instance + + monkeypatch.setattr("autogen_ext.tools.mcp._workbench.McpSessionActor", mock_actor_class) + + # Create workbench without model client + workbench = McpWorkbench(server_params=sample_server_params) + + # Mock list_tools response + list_tools_result = ListToolsResult(tools=sample_mcp_tools) + future_result: asyncio.Future[ListToolsResult] = asyncio.Future() + future_result.set_result(list_tools_result) + mock_mcp_actor.call.return_value = future_result + + try: + # Start the workbench - this should create the actor without model client + await workbench.start() + + # Verify that the actor was created without the model client + mock_actor_class.assert_called_once_with(sample_server_params, model_client=None) + + # Verify that initialize was called on the actor + mock_mcp_actor.initialize.assert_called_once() + + finally: + if workbench._actor: # type: ignore[attr-defined] + workbench._actor = None # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_mcp_workbench_list_tools_with_model_client( + sample_mcp_tools: list[McpTool], + mock_mcp_actor: AsyncMock, + sample_server_params: StdioServerParams, + mock_model_client: MockChatCompletionClient, +) -> None: + """Test that list_tools works correctly with model client.""" + # Create workbench with model client + workbench = McpWorkbench( + server_params=sample_server_params, + model_client=mock_model_client, + ) + workbench._actor = mock_mcp_actor # type: ignore[attr-defined] + + # Mock list_tools response + list_tools_result = ListToolsResult(tools=sample_mcp_tools) + future_result: asyncio.Future[ListToolsResult] = asyncio.Future() + future_result.set_result(list_tools_result) + mock_mcp_actor.call.return_value = future_result + + try: + # List tools + tools = await workbench.list_tools() + + # Verify tools are returned correctly + assert len(tools) == 1 + assert tools[0]["name"] == "test_tool" + assert tools[0].get("description") == "A test tool" + + # Verify actor was called correctly + mock_mcp_actor.call.assert_called_with("list_tools", None) + + finally: + workbench._actor = None # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_mcp_workbench_call_tool_with_model_client( + mock_mcp_actor: AsyncMock, + sample_server_params: StdioServerParams, + mock_model_client: MockChatCompletionClient, +) -> None: + """Test that call_tool works correctly with model client.""" + # Create workbench with model client + workbench = McpWorkbench( + server_params=sample_server_params, + model_client=mock_model_client, + ) + workbench._actor = mock_mcp_actor # type: ignore[attr-defined] + + # Mock call_tool response + call_tool_result = CallToolResult( + content=[TextContent(type="text", text="Tool executed successfully")], + isError=False, + ) + future_result: asyncio.Future[CallToolResult] = asyncio.Future() + future_result.set_result(call_tool_result) + mock_mcp_actor.call.return_value = future_result + + try: + # Call tool + result = await workbench.call_tool( + name="test_tool", + arguments={"param": "test_value"}, + ) + + # Verify result + assert not result.is_error + assert len(result.result) == 1 + assert result.result[0].content == "Tool executed successfully" + + # Verify actor was called correctly + mock_mcp_actor.call.assert_called_with("call_tool", {"name": "test_tool", "kargs": {"param": "test_value"}}) + + finally: + workbench._actor = None # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_mcp_workbench_state_operations_with_model_client( + sample_server_params: StdioServerParams, + mock_model_client: MockChatCompletionClient, +) -> None: + """Test state save/load operations with model client.""" + # Create workbench with model client + workbench = McpWorkbench( + server_params=sample_server_params, + model_client=mock_model_client, + ) + + # Save state + state = await workbench.save_state() + assert state["type"] == "McpWorkBenchState" + + # Load state (should not affect model client) + await workbench.load_state(state) + + # Model client should still be there + assert workbench._model_client is mock_model_client # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_mcp_workbench_fallback_start_behavior( + sample_mcp_tools: list[McpTool], + mock_mcp_actor: AsyncMock, + sample_server_params: StdioServerParams, + mock_model_client: MockChatCompletionClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test fallback start behavior when actor is not initialized.""" + # Mock the McpSessionActor creation + mock_actor_class = MagicMock() + mock_actor_instance = mock_mcp_actor + mock_actor_class.return_value = mock_actor_instance + + monkeypatch.setattr("autogen_ext.tools.mcp._workbench.McpSessionActor", mock_actor_class) + + # Create workbench with model client + workbench = McpWorkbench( + server_params=sample_server_params, + model_client=mock_model_client, + ) + + # Ensure actor is not initialized + assert workbench._actor is None # type: ignore[attr-defined] + + # Mock list_tools response + list_tools_result = ListToolsResult(tools=sample_mcp_tools) + future_result: asyncio.Future[ListToolsResult] = asyncio.Future() + future_result.set_result(list_tools_result) + mock_mcp_actor.call.return_value = future_result + + try: + # Call list_tools without explicitly starting - should auto-start + tools = await workbench.list_tools() + + # Verify that the actor was created with the model client + mock_actor_class.assert_called_once_with(sample_server_params, model_client=mock_model_client) + + # Verify that initialize was called on the actor + mock_mcp_actor.initialize.assert_called_once() + + # Verify tools are returned correctly + assert len(tools) == 1 + assert tools[0]["name"] == "test_tool" + + finally: + if workbench._actor: # type: ignore[attr-defined] + workbench._actor = None # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_mcp_workbench_model_client_vision_support( + sample_server_params: StdioServerParams, + mock_model_client_with_vision: MockChatCompletionClient, +) -> None: + """Test that model client with vision support is properly configured.""" + # Create workbench with vision-enabled model client + workbench = McpWorkbench( + server_params=sample_server_params, + model_client=mock_model_client_with_vision, + ) + + # Verify model client has vision capabilities + assert workbench._model_client is not None # type: ignore[attr-defined] + assert workbench._model_client.model_info["vision"] is True # type: ignore[attr-defined] + assert workbench._model_client.model_info["family"] == "test-vision-model" # type: ignore[attr-defined] + + # Test serialization preserves vision capabilities + config = workbench._to_config() # type: ignore[attr-defined] + loaded_workbench = McpWorkbench._from_config(config) # type: ignore[attr-defined] + + assert loaded_workbench._model_client is not None # type: ignore[attr-defined] + assert loaded_workbench._model_client.model_info["vision"] is True # type: ignore[attr-defined] + assert loaded_workbench._model_client.model_info["family"] == "test-vision-model" # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_mcp_workbench_component_interface( + sample_server_params: StdioServerParams, + mock_model_client: MockChatCompletionClient, +) -> None: + """Test that McpWorkbench properly implements the Component interface.""" + # Test component provider override + assert McpWorkbench.component_provider_override == "autogen_ext.tools.mcp.McpWorkbench" + + # Test config schema + config_schema = McpWorkbench.component_config_schema + assert config_schema is not None + + # Test serialization/deserialization through component interface + workbench = McpWorkbench( + server_params=sample_server_params, + model_client=mock_model_client, + ) + + # Test _to_config + config = workbench._to_config() # type: ignore[attr-defined] + assert hasattr(config, "server_params") + assert hasattr(config, "model_client") + assert hasattr(config, "tool_overrides") + + # Test _from_config + loaded_workbench = McpWorkbench._from_config(config) # type: ignore[attr-defined] + assert loaded_workbench.server_params == sample_server_params + assert loaded_workbench._model_client is not None # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_sampling_callback_invocation( + mock_model_client: MockChatCompletionClient, sample_server_params: StdioServerParams +) -> None: + """Test that the sampling_callback in McpSessionActor is actually called when requested.""" + from unittest.mock import MagicMock + + from autogen_ext.tools.mcp._actor import McpSessionActor + from mcp.shared.context import RequestContext + from mcp.types import CreateMessageRequestParams, CreateMessageResult, SamplingMessage, TextContent + + # Set up the mock for this specific test + mock_create = AsyncMock( + return_value=CreateResult( + content="Mock response", + finish_reason="stop", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ) + mock_model_client.create = mock_create # type: ignore[method-assign] + + # Create an actor with a model client + actor = McpSessionActor(sample_server_params, model_client=mock_model_client) + + # Create mock request context and parameters for sampling + mock_context = MagicMock(spec=RequestContext) + sample_params = CreateMessageRequestParams( + messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello, how are you?"))], + maxTokens=100, + ) + + # Test the sampling callback directly + result = await actor._sampling_callback(mock_context, sample_params) # type: ignore[attr-defined] + + # Verify the result is a CreateMessageResult (not an error) + assert isinstance(result, CreateMessageResult) + assert result.role == "assistant" + assert result.content.type == "text" + assert result.content.text == "Mock response" + assert result.model == "test-model" + assert result.stopReason == "stop" + + # Verify the model client was called correctly + assert mock_create.called + # Get the call arguments + call_args = mock_create.call_args + messages = call_args.kwargs["messages"] # Messages are passed as keyword argument + assert len(messages) == 1 + assert messages[0].source == "user" + assert "Hello, how are you?" in str(messages[0].content) + + +@pytest.mark.asyncio +async def test_sampling_callback_without_model_client(sample_server_params: StdioServerParams) -> None: + """Test that sampling_callback returns error when no model client is provided.""" + from unittest.mock import MagicMock + + from autogen_ext.tools.mcp._actor import McpSessionActor + from mcp.shared.context import RequestContext + from mcp.types import CreateMessageRequestParams, ErrorData, SamplingMessage, TextContent + + # Create an actor without a model client + actor = McpSessionActor(sample_server_params, model_client=None) + + # Create mock request context and parameters + mock_context = MagicMock(spec=RequestContext) + sample_params = CreateMessageRequestParams( + messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello, how are you?"))], + maxTokens=100, + ) + + # Test the sampling callback + result = await actor._sampling_callback(mock_context, sample_params) # type: ignore[attr-defined] + + # Verify the result is an ErrorData (indicating no model client) + assert isinstance(result, ErrorData) + + +@pytest.mark.asyncio +async def test_sampling_callback_with_system_prompt( + mock_model_client: MockChatCompletionClient, sample_server_params: StdioServerParams +) -> None: + """Test sampling_callback with system prompt.""" + from unittest.mock import MagicMock + + from autogen_ext.tools.mcp._actor import McpSessionActor + from mcp.shared.context import RequestContext + from mcp.types import CreateMessageRequestParams, CreateMessageResult, SamplingMessage, TextContent + + # Set up the mock for this specific test + mock_create = AsyncMock( + return_value=CreateResult( + content="Mock response", + finish_reason="stop", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ) + mock_model_client.create = mock_create # type: ignore[method-assign] + + # Create an actor with a model client + actor = McpSessionActor(sample_server_params, model_client=mock_model_client) + + # Create mock request context and parameters with system prompt + mock_context = MagicMock(spec=RequestContext) + sample_params = CreateMessageRequestParams( + messages=[SamplingMessage(role="user", content=TextContent(type="text", text="What's the weather?"))], + systemPrompt="You are a helpful weather assistant.", + maxTokens=100, + ) + + # Test the sampling callback + result = await actor._sampling_callback(mock_context, sample_params) # type: ignore[attr-defined] + + # Verify the result + assert isinstance(result, CreateMessageResult) + + # Verify the model client was called with system prompt + assert mock_create.called + call_args = mock_create.call_args + messages = call_args.kwargs["messages"] + assert len(messages) == 2 # system + user message + + # Check system message + system_msg = messages[0] + assert system_msg.__class__.__name__ == "SystemMessage" + assert "helpful weather assistant" in str(system_msg.content) + + # Check user message + user_msg = messages[1] + assert user_msg.source == "user" + assert "weather" in str(user_msg.content) + + +@pytest.mark.asyncio +async def test_sampling_callback_with_image_content( + mock_model_client_with_vision: MockChatCompletionClient, sample_server_params: StdioServerParams +) -> None: + """Test sampling_callback with image content when vision is supported.""" + from unittest.mock import MagicMock + + from autogen_ext.tools.mcp._actor import McpSessionActor + from mcp.shared.context import RequestContext + from mcp.types import CreateMessageRequestParams, CreateMessageResult, ImageContent, SamplingMessage + + # Set up the mock for this specific test + mock_create = AsyncMock( + return_value=CreateResult( + content="Mock response", + finish_reason="stop", + usage=RequestUsage(prompt_tokens=10, completion_tokens=5), + cached=False, + ) + ) + mock_model_client_with_vision.create = mock_create # type: ignore[method-assign] + + # Create an actor with a vision-capable model client + actor = McpSessionActor(sample_server_params, model_client=mock_model_client_with_vision) + + # Create a simple test image (1x1 pixel PNG) - using a proper PNG + # This is a valid 1x1 transparent PNG + test_image_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" + + # Create mock request context and parameters with image + mock_context = MagicMock(spec=RequestContext) + sample_params = CreateMessageRequestParams( + messages=[ + SamplingMessage(role="user", content=ImageContent(type="image", data=test_image_b64, mimeType="image/png")) + ], + maxTokens=100, + ) + + # Test the sampling callback + result = await actor._sampling_callback(mock_context, sample_params) # type: ignore[attr-defined] + + # Verify the result + assert isinstance(result, CreateMessageResult) + + # Verify the model client was called + assert mock_create.called + call_args = mock_create.call_args + messages = call_args.kwargs["messages"] + assert len(messages) == 1 + + # Check that the message contains image content + user_msg = messages[0] + assert user_msg.source == "user" + # The content should be a list containing an Image object + assert isinstance(user_msg.content, list) + assert len(user_msg.content) == 1 # type: ignore[arg-type] + + +@pytest.mark.asyncio +async def test_sampling_callback_image_without_vision_support( + mock_model_client: MockChatCompletionClient, sample_server_params: StdioServerParams +) -> None: + """Test sampling_callback fails gracefully when image is sent to non-vision model.""" + from unittest.mock import MagicMock + + from autogen_ext.tools.mcp._actor import McpSessionActor + from mcp.shared.context import RequestContext + from mcp.types import CreateMessageRequestParams, ErrorData, ImageContent, SamplingMessage + + # Create an actor with a non-vision model client + actor = McpSessionActor(sample_server_params, model_client=mock_model_client) + + # Create a simple test image - using the same valid PNG + test_image_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" + + # Create mock request context and parameters with image + mock_context = MagicMock(spec=RequestContext) + sample_params = CreateMessageRequestParams( + messages=[ + SamplingMessage(role="user", content=ImageContent(type="image", data=test_image_b64, mimeType="image/png")) + ], + maxTokens=100, + ) + + # Test the sampling callback + result = await actor._sampling_callback(mock_context, sample_params) # type: ignore[attr-defined] + + # Verify the result is an error due to vision not being supported + assert isinstance(result, ErrorData) + # Check both the message and data fields for vision/image-related error + error_text = (result.message + " " + str(result.data or "")).lower() + assert "vision" in error_text or "image" in error_text + + +@pytest.mark.asyncio +async def test_sampling_callback_model_client_error(sample_server_params: StdioServerParams) -> None: + """Test sampling_callback handles model client errors gracefully.""" + from unittest.mock import AsyncMock, MagicMock + + from autogen_ext.tools.mcp._actor import McpSessionActor + from mcp.shared.context import RequestContext + from mcp.types import CreateMessageRequestParams, ErrorData, SamplingMessage, TextContent + + # Create a model client that will fail + failing_model_client = MockChatCompletionClient() + failing_model_client.create = AsyncMock(side_effect=Exception("Model API error")) # type: ignore[method-assign] + + # Create an actor with the failing model client + actor = McpSessionActor(sample_server_params, model_client=failing_model_client) + + # Create mock request context and parameters + mock_context = MagicMock(spec=RequestContext) + sample_params = CreateMessageRequestParams( + messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))], maxTokens=100 + ) + + # Test the sampling callback + result = await actor._sampling_callback(mock_context, sample_params) # type: ignore[attr-defined] + + # Verify the result is an error + assert isinstance(result, ErrorData) + assert result.data is not None and "Model API error" in str(result.data) + + +@pytest.mark.asyncio +async def test_sampling_callback_integration_through_session(sample_server_params: StdioServerParams) -> None: + """Test that shows how sampling_callback should work when properly integrated with session.""" + import asyncio + from unittest.mock import AsyncMock, MagicMock, patch + + # Create a mock model client + mock_model_client = MagicMock() + mock_model_client.model_info = {"family": "test", "vision": False} + mock_model_client.create = AsyncMock(return_value=MagicMock(content="Mock response", finish_reason="stop")) + + # Mock the session creation to verify the sampling_callback is passed + with patch("autogen_ext.tools.mcp._actor.create_mcp_server_session") as mock_session_factory: + # Create mock session that properly handles the context manager protocol + mock_session = AsyncMock() + mock_session.initialize = AsyncMock() + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=None) + mock_session_factory.return_value = mock_session + + from autogen_ext.tools.mcp._actor import McpSessionActor + + actor = McpSessionActor(sample_server_params, model_client=mock_model_client) + + # Initialize the actor (this will call create_mcp_server_session) + await actor.initialize() + + # Wait a moment for the actor to initialize + await asyncio.sleep(0.1) + + # Verify that the session was created with the sampling_callback + mock_session_factory.assert_called_once_with(sample_server_params, sampling_callback=actor._sampling_callback) # type: ignore[attr-defined] + + # Clean up by trying to close the actor gracefully + if actor._active and actor._actor_task is not None: # type: ignore[attr-defined] + try: + await actor.close() + except Exception: + # Ignore cleanup errors since the main assertion passed + pass diff --git a/python/packages/autogen-ext/tests/tools/test_mcp_workbench_overrides.py b/python/packages/autogen-ext/tests/tools/test_mcp_workbench_overrides.py new file mode 100644 index 000000000000..76fcef1c13ed --- /dev/null +++ b/python/packages/autogen-ext/tests/tools/test_mcp_workbench_overrides.py @@ -0,0 +1,298 @@ +import asyncio +from typing import Any, Dict +from unittest.mock import AsyncMock + +import pytest +from autogen_core.tools import ToolOverride +from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams +from mcp import Tool +from mcp.types import ListToolsResult + + +@pytest.fixture +def sample_mcp_tools() -> list[Tool]: + """Create sample MCP tools for testing.""" + return [ + Tool( + name="fetch", + description="Fetches content from a URL", + inputSchema={ + "type": "object", + "properties": {"url": {"type": "string"}}, + "required": ["url"], + }, + ), + Tool( + name="search", + description="Searches for information", + inputSchema={ + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + ), + ] + + +@pytest.fixture +def mock_mcp_actor() -> AsyncMock: + """Mock MCP session actor.""" + actor = AsyncMock() + return actor + + +@pytest.fixture +def sample_server_params() -> StdioServerParams: + """Sample server parameters for testing.""" + return StdioServerParams(command="echo", args=["test"]) + + +@pytest.mark.asyncio +async def test_mcp_workbench_with_tool_overrides( + sample_mcp_tools: list[Tool], mock_mcp_actor: AsyncMock, sample_server_params: StdioServerParams +) -> None: + """Test McpWorkbench with tool name and description overrides.""" + + # Define tool overrides + overrides: Dict[str, ToolOverride] = { + "fetch": ToolOverride(name="web_fetch", description="Enhanced web fetching tool"), + "search": ToolOverride(description="Advanced search functionality"), # Only override description + } + + # Create workbench with overrides + workbench = McpWorkbench(server_params=sample_server_params, tool_overrides=overrides) + workbench._actor = mock_mcp_actor # type: ignore[reportPrivateUsage] + + # Mock list_tools response + list_tools_result = ListToolsResult(tools=sample_mcp_tools) + + # The actor.call() should return a Future that when awaited returns the list_tools_result + future_result: asyncio.Future[ListToolsResult] = asyncio.Future() + future_result.set_result(list_tools_result) + mock_mcp_actor.call.return_value = future_result + + try: + # List tools and verify overrides are applied + tools = await workbench.list_tools() + assert len(tools) == 2 + + # Check first tool has name and description overridden + assert tools[0].get("name") == "web_fetch" + assert tools[0].get("description") == "Enhanced web fetching tool" + + # Check second tool has only description overridden + assert tools[1].get("name") == "search" # Original name + assert tools[1].get("description") == "Advanced search functionality" # Overridden description + + # Verify actor was called correctly + mock_mcp_actor.call.assert_called_with("list_tools", None) + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_mcp_workbench_call_tool_with_overrides( + sample_mcp_tools: list[Tool], mock_mcp_actor: AsyncMock, sample_server_params: StdioServerParams +) -> None: + """Test calling tools with override names maps back to original names.""" + + overrides: Dict[str, ToolOverride] = { + "fetch": ToolOverride(name="web_fetch", description="Enhanced web fetching tool") + } + + workbench = McpWorkbench(server_params=sample_server_params, tool_overrides=overrides) + workbench._actor = mock_mcp_actor # type: ignore[reportPrivateUsage] + + # Mock successful tool call response + from mcp.types import CallToolResult, TextContent + + mock_result = CallToolResult(content=[TextContent(text="Mock response", type="text")], isError=False) + + # Create futures for each call + def mock_call_side_effect( + method: str, args: dict[str, Any] | None = None + ) -> asyncio.Future[ListToolsResult | CallToolResult]: + future_result: asyncio.Future[ListToolsResult | CallToolResult] = asyncio.Future() + if method == "list_tools": + future_result.set_result(ListToolsResult(tools=sample_mcp_tools)) + elif method == "call_tool": + future_result.set_result(mock_result) + else: + future_result.set_exception(ValueError(f"Unexpected method: {method}")) + return future_result + + mock_mcp_actor.call.side_effect = mock_call_side_effect + + try: + # Call tool using override name + result = await workbench.call_tool("web_fetch", {"url": "https://example.com"}) + + # Verify the result + assert result.name == "web_fetch" # Should return the override name + assert result.result[0].content == "Mock response" + assert result.is_error is False + + # Verify the actor was called with the original tool name + call_args = mock_mcp_actor.call.call_args_list[-1] + assert call_args[0][0] == "call_tool" + assert call_args[0][1]["name"] == "fetch" # Original name should be used + assert call_args[0][1]["kargs"] == {"url": "https://example.com"} + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_mcp_workbench_without_overrides( + sample_mcp_tools: list[Tool], mock_mcp_actor: AsyncMock, sample_server_params: StdioServerParams +) -> None: + """Test McpWorkbench without overrides (original behavior).""" + + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = mock_mcp_actor # type: ignore[reportPrivateUsage] + + # Mock list_tools response + list_tools_result = ListToolsResult(tools=sample_mcp_tools) + future_result: asyncio.Future[ListToolsResult] = asyncio.Future() + future_result.set_result(list_tools_result) + mock_mcp_actor.call.return_value = future_result + + try: + tools = await workbench.list_tools() + assert len(tools) == 2 + + # Verify original names and descriptions are preserved + assert tools[0].get("name") == "fetch" + assert tools[0].get("description") == "Fetches content from a URL" + assert tools[1].get("name") == "search" + assert tools[1].get("description") == "Searches for information" + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_mcp_workbench_serialization_with_overrides(sample_server_params: StdioServerParams) -> None: + """Test that McpWorkbench can be serialized and deserialized with overrides.""" + + overrides: Dict[str, ToolOverride] = { + "fetch": ToolOverride(name="web_fetch", description="Enhanced web fetching tool") + } + + # Create workbench with overrides + workbench = McpWorkbench(server_params=sample_server_params, tool_overrides=overrides) + + # Save configuration + config = workbench.dump_component() + assert "tool_overrides" in config.config + assert "fetch" in config.config["tool_overrides"] + assert config.config["tool_overrides"]["fetch"]["name"] == "web_fetch" + assert config.config["tool_overrides"]["fetch"]["description"] == "Enhanced web fetching tool" + + # Load workbench from configuration + new_workbench = McpWorkbench.load_component(config) + assert len(new_workbench._tool_overrides) == 1 # type: ignore[reportPrivateUsage] + assert new_workbench._tool_overrides["fetch"].name == "web_fetch" # type: ignore[reportPrivateUsage] + assert new_workbench._tool_overrides["fetch"].description == "Enhanced web fetching tool" # type: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_mcp_workbench_partial_overrides( + sample_mcp_tools: list[Tool], mock_mcp_actor: AsyncMock, sample_server_params: StdioServerParams +) -> None: + """Test McpWorkbench with partial overrides (name only, description only).""" + + overrides: Dict[str, ToolOverride] = { + "fetch": ToolOverride(name="web_fetch"), # Only name override + "search": ToolOverride(description="Advanced search"), # Only description override + } + + workbench = McpWorkbench(server_params=sample_server_params, tool_overrides=overrides) + workbench._actor = mock_mcp_actor # type: ignore[reportPrivateUsage] + + # Mock list_tools response + list_tools_result = ListToolsResult(tools=sample_mcp_tools) + future_result: asyncio.Future[ListToolsResult] = asyncio.Future() + future_result.set_result(list_tools_result) + mock_mcp_actor.call.return_value = future_result + + try: + tools = await workbench.list_tools() + + # fetch: name overridden, description unchanged + assert tools[0].get("name") == "web_fetch" + assert tools[0].get("description") == "Fetches content from a URL" # Original description + + # search: name unchanged, description overridden + assert tools[1].get("name") == "search" # Original name + assert tools[1].get("description") == "Advanced search" # Overridden description + + finally: + workbench._actor = None # type: ignore[reportPrivateUsage] + + +def test_mcp_tool_override_model() -> None: + """Test ToolOverride model functionality for MCP.""" + + # Test with both fields + override1 = ToolOverride(name="new_name", description="new_desc") + assert override1.name == "new_name" + assert override1.description == "new_desc" + + # Test with only name + override2 = ToolOverride(name="new_name") + assert override2.name == "new_name" + assert override2.description is None + + # Test with only description + override3 = ToolOverride(description="new_desc") + assert override3.name is None + assert override3.description == "new_desc" + + # Test empty + override4 = ToolOverride() + assert override4.name is None + assert override4.description is None + + +@pytest.mark.asyncio +async def test_mcp_workbench_override_name_to_original_mapping(sample_server_params: StdioServerParams) -> None: + """Test that the reverse mapping from override names to original names works correctly.""" + + overrides: Dict[str, ToolOverride] = { + "original1": ToolOverride(name="override1"), + "original2": ToolOverride(name="override2"), + "original3": ToolOverride(description="only description override"), # No name change, only description + } + + workbench = McpWorkbench(server_params=sample_server_params, tool_overrides=overrides) + + # Check reverse mapping is built correctly + assert workbench._override_name_to_original["override1"] == "original1" # type: ignore[reportPrivateUsage] + assert workbench._override_name_to_original["override2"] == "original2" # type: ignore[reportPrivateUsage] + assert "original3" not in workbench._override_name_to_original # type: ignore[reportPrivateUsage] + assert len(workbench._override_name_to_original) == 2 # type: ignore[reportPrivateUsage] + + +def test_mcp_workbench_conflict_detection() -> None: + """Test that McpWorkbench detects conflicts in tool override names.""" + + server_params = StdioServerParams(command="echo", args=["test"]) + + # Test 1: Valid overrides - should work + overrides_valid: Dict[str, ToolOverride] = { + "fetch": ToolOverride(name="web_fetch"), + "search": ToolOverride(name="advanced_search"), + } + workbench_valid = McpWorkbench(server_params=server_params, tool_overrides=overrides_valid) + assert workbench_valid._override_name_to_original["web_fetch"] == "fetch" # type: ignore[reportPrivateUsage] + assert workbench_valid._override_name_to_original["advanced_search"] == "search" # type: ignore[reportPrivateUsage] + + # Test 2: Duplicate override names - should fail + overrides_duplicate: Dict[str, ToolOverride] = { + "fetch": ToolOverride(name="same_name"), + "search": ToolOverride(name="same_name"), # Duplicate + } + with pytest.raises(ValueError): + McpWorkbench(server_params=server_params, tool_overrides=overrides_duplicate) diff --git a/python/packages/autogen-ext/tests/tools/test_mcp_workbench_warnings_and_errors.py b/python/packages/autogen-ext/tests/tools/test_mcp_workbench_warnings_and_errors.py new file mode 100644 index 000000000000..ea598f4e55b3 --- /dev/null +++ b/python/packages/autogen-ext/tests/tools/test_mcp_workbench_warnings_and_errors.py @@ -0,0 +1,305 @@ +"""Tests for McpWorkbench expected errors and warnings.""" + +import asyncio +import builtins +from typing import Any +from unittest.mock import AsyncMock, patch + +import pytest +from autogen_core import Image +from autogen_core.tools import ImageResultContent, TextResultContent +from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams +from mcp.types import ( + CallToolResult, + EmbeddedResource, + ImageContent, + TextContent, +) + + +@pytest.fixture +def sample_server_params() -> StdioServerParams: + """Sample server parameters for testing.""" + return StdioServerParams(command="echo", args=["test"]) + + +@pytest.fixture +def mock_actor() -> AsyncMock: + """Mock actor for testing.""" + return AsyncMock() + + +@pytest.mark.asyncio +async def test_list_tools_actor_none_after_start(sample_server_params: StdioServerParams) -> None: + """Test list_tools when actor is None after start attempt - covers line 274.""" + workbench = McpWorkbench(server_params=sample_server_params) + + # Mock start method to set _actor to None + with patch.object(workbench, "start") as mock_start: + mock_start.return_value = None + workbench._actor = None # type: ignore[reportPrivateUsage] + + with pytest.raises(RuntimeError, match="Actor is not initialized. Please check the server connection."): + await workbench.list_tools() + + +@pytest.mark.asyncio +async def test_call_tool_actor_none_after_start(sample_server_params: StdioServerParams) -> None: + """Test call_tool when actor is None after start attempt - covers line 320.""" + workbench = McpWorkbench(server_params=sample_server_params) + + # Mock start method to set _actor to None + with patch.object(workbench, "start") as mock_start: + mock_start.return_value = None + workbench._actor = None # type: ignore[reportPrivateUsage] + + with pytest.raises(RuntimeError, match="Actor is not initialized. Please check the server connection."): + await workbench.call_tool("test_tool") + + +@pytest.mark.asyncio +async def test_list_prompts_actor_none_after_start(sample_server_params: StdioServerParams) -> None: + """Test list_prompts when actor is None after start attempt - covers line 364.""" + workbench = McpWorkbench(server_params=sample_server_params) + + # Mock start method to set _actor to None + with patch.object(workbench, "start") as mock_start: + mock_start.return_value = None + workbench._actor = None # type: ignore[reportPrivateUsage] + + with pytest.raises(RuntimeError, match="Actor is not initialized. Please check the server connection."): + await workbench.list_prompts() + + +@pytest.mark.asyncio +async def test_call_tool_image_content_handling(sample_server_params: StdioServerParams, mock_actor: AsyncMock) -> None: + """Test call_tool with ImageContent - covers lines 346-347.""" + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = mock_actor # type: ignore[reportPrivateUsage] + + # Mock tool result with ImageContent + image_content = ImageContent( + type="image", + data="iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==", + mimeType="image/png", + ) + + call_result = CallToolResult(content=[image_content], isError=False) + + # Mock the call method to return a future that resolves to the result + future: asyncio.Future[CallToolResult] = asyncio.Future() + future.set_result(call_result) + mock_actor.call.return_value = future + + result = await workbench.call_tool("test_tool") + + assert len(result.result) == 1 + assert isinstance(result.result[0], ImageResultContent) + assert isinstance(result.result[0].content, Image) + + +@pytest.mark.asyncio +async def test_call_tool_embedded_resource_handling( + sample_server_params: StdioServerParams, mock_actor: AsyncMock +) -> None: + """Test call_tool with EmbeddedResource - covers lines 348-351.""" + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = mock_actor # type: ignore[reportPrivateUsage] + + # Mock tool result with EmbeddedResource + from mcp.types import TextResourceContents + from pydantic import AnyUrl + + embedded_resource = EmbeddedResource( + type="resource", + resource=TextResourceContents(uri=AnyUrl("test://resource"), text="test content"), + ) + + call_result = CallToolResult(content=[embedded_resource], isError=False) + + # Mock the call method to return a future that resolves to the result + future: asyncio.Future[CallToolResult] = asyncio.Future() + future.set_result(call_result) + mock_actor.call.return_value = future + + result = await workbench.call_tool("test_tool") + + assert len(result.result) == 1 + assert isinstance(result.result[0], TextResultContent) + # Should contain JSON representation of the embedded resource + assert "resource" in result.result[0].content + + +@pytest.mark.asyncio +async def test_call_tool_exception_handling(sample_server_params: StdioServerParams, mock_actor: AsyncMock) -> None: + """Test call_tool exception handling - covers lines 354-357.""" + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = mock_actor # type: ignore[reportPrivateUsage] + + # Mock actor to raise an exception + mock_actor.call.side_effect = Exception("Test exception") + + result = await workbench.call_tool("test_tool") + + assert result.is_error + assert len(result.result) == 1 + assert isinstance(result.result[0], TextResultContent) + assert "Test exception" in result.result[0].content + + +def test_format_errors_with_exception_group() -> None: + """Test _format_errors with ExceptionGroup - covers lines 444-452.""" + workbench = McpWorkbench(server_params=StdioServerParams(command="echo", args=["test"])) + + # Only test if ExceptionGroup is available (Python 3.11+) + if hasattr(builtins, "ExceptionGroup"): + # Create an ExceptionGroup + sub_exceptions = [ + ValueError("Error 1"), + RuntimeError("Error 2"), + ] + exception_group = builtins.ExceptionGroup("Multiple errors", sub_exceptions) + + result = workbench._format_errors(exception_group) # type: ignore[reportPrivateUsage] + + assert "Error 1" in result + assert "Error 2" in result + else: + # For Python < 3.11, just test regular exception handling + regular_exception = ValueError("Regular error") + result = workbench._format_errors(regular_exception) # type: ignore[reportPrivateUsage] + assert "Regular error" in result + + +@pytest.mark.asyncio +async def test_call_tool_with_none_arguments(sample_server_params: StdioServerParams, mock_actor: AsyncMock) -> None: + """Test call_tool with None arguments - covers line 323.""" + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = mock_actor # type: ignore[reportPrivateUsage] + + # Mock successful tool call + call_result = CallToolResult(content=[TextContent(type="text", text="Success")], isError=False) + + # Mock the call method to return a future that resolves to the result + future: asyncio.Future[CallToolResult] = asyncio.Future() + future.set_result(call_result) + mock_actor.call.return_value = future + + # Call with None arguments + result = await workbench.call_tool("test_tool", arguments=None) + + # Should handle None arguments gracefully + assert not result.is_error + assert len(result.result) == 1 + + +@pytest.mark.asyncio +async def test_call_tool_with_none_cancellation_token( + sample_server_params: StdioServerParams, mock_actor: AsyncMock +) -> None: + """Test call_tool with None cancellation_token - covers line 322.""" + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = mock_actor # type: ignore[reportPrivateUsage] + + # Mock successful tool call + call_result = CallToolResult(content=[TextContent(type="text", text="Success")], isError=False) + + # Mock the call method to return a future that resolves to the result + future: asyncio.Future[CallToolResult] = asyncio.Future() + future.set_result(call_result) + mock_actor.call.return_value = future + + # Call with None cancellation_token + result = await workbench.call_tool("test_tool", cancellation_token=None) + + # Should handle None cancellation_token gracefully + assert not result.is_error + assert len(result.result) == 1 + + +@pytest.mark.asyncio +async def test_initialize_result_property_with_actor( + sample_server_params: StdioServerParams, mock_actor: AsyncMock +) -> None: + """Test initialize_result property when actor exists.""" + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = mock_actor # type: ignore[reportPrivateUsage] + mock_actor.initialize_result = "test_result" + + result = workbench.initialize_result # type: ignore[reportPrivateUsage] + assert result == "test_result" + + +@pytest.mark.asyncio +async def test_initialize_result_property_without_actor(sample_server_params: StdioServerParams) -> None: + """Test initialize_result property when actor is None.""" + workbench = McpWorkbench(server_params=sample_server_params) + workbench._actor = None # type: ignore[reportPrivateUsage] + + result = workbench.initialize_result # type: ignore[reportPrivateUsage] + assert result is None + + +@pytest.mark.asyncio +async def test_to_config_method(sample_server_params: StdioServerParams) -> None: + """Test _to_config method.""" + workbench = McpWorkbench(server_params=sample_server_params) + config = workbench._to_config() # type: ignore[reportPrivateUsage] + assert config.server_params == sample_server_params + + +@pytest.mark.asyncio +async def test_from_config_method(sample_server_params: StdioServerParams) -> None: + """Test _from_config method.""" + from autogen_ext.tools.mcp._workbench import McpWorkbenchConfig # type: ignore[reportPrivateUsage] + + config = McpWorkbenchConfig(server_params=sample_server_params) + workbench = McpWorkbench._from_config(config) # type: ignore[reportPrivateUsage] + assert workbench.server_params == sample_server_params + + +@pytest.mark.asyncio +async def test_async_context_manager(sample_server_params: StdioServerParams) -> None: + """Test async context manager functionality.""" + workbench = McpWorkbench(server_params=sample_server_params) + + # Test that the context manager properly handles start/stop + with patch.object(workbench, "start") as mock_start, patch.object(workbench, "stop") as mock_stop: + async with workbench: + pass + + mock_start.assert_called_once() + mock_stop.assert_called_once() + + +@pytest.mark.asyncio +async def test_sampling_callback_functionality(sample_server_params: StdioServerParams) -> None: + """Test sampling callback functionality for private method access.""" + workbench = McpWorkbench(server_params=sample_server_params) + + # Create a mock that simulates the sampling callback + mock_callback: AsyncMock = AsyncMock() + + # Test that the workbench can handle sampling callbacks + # This tests private method access which appears in the error reports + workbench._sampling_callback = mock_callback # type: ignore[attr-defined] + + # Verify the callback was set + assert workbench._sampling_callback is mock_callback # type: ignore[attr-defined] + + +def test_misc_lambda_types() -> None: + """Test miscellaneous lambda types for coverage.""" + + # Test lambda with unknown parameter types + def test_lambda(obj: Any, cls: Any) -> bool: + return True + + assert test_lambda("test", str) is True + + # Test lambda with return type + def test_lambda2(x: Any) -> bool: + return isinstance(x, str) + + assert test_lambda2("test") is True + assert test_lambda2(123) is False diff --git a/python/packages/autogen-studio/README.md b/python/packages/autogen-studio/README.md index 75393d467a8a..8f7452de77fa 100644 --- a/python/packages/autogen-studio/README.md +++ b/python/packages/autogen-studio/README.md @@ -124,6 +124,45 @@ AutoGen Studio also takes several parameters to customize the application: Now that you have AutoGen Studio installed and running, you are ready to explore its capabilities, including defining and modifying agent workflows, interacting with agents and sessions, and expanding agent skills. +## AutoGen Studio Lite + +AutoGen Studio Lite provides a lightweight way to quickly prototype and experiment with AI agent teams. It's designed for rapid experimentation without the full database setup. + +### CLI Usage + +Launch Studio Lite from the command line: + +```bash +# Quick start with default team +autogenstudio lite + +# Use custom team file +autogenstudio lite --team ./my_team.json --port 8080 + +# Custom session name +autogenstudio lite --session-name "My Experiment" --auto-open +``` + +### Programmatic Usage + +Use Studio Lite directly in your Python code: + +```python +from autogenstudio.lite import LiteStudio + +# Quick start with default team +studio = LiteStudio() +# Use with AutoGen team objects +from autogen_agentchat.teams import RoundRobinGroupChat +team = RoundRobinGroupChat([agent1, agent2], termination_condition=...) + +# Context manager usage +with LiteStudio(team=team) as studio: + # Studio runs in background + # Do other work here + pass +``` + #### Local frontend development server See `./frontend/README.md` diff --git a/python/packages/autogen-studio/autogenstudio/cli.py b/python/packages/autogen-studio/autogenstudio/cli.py index 8bf27915c1a4..373f13e087f2 100644 --- a/python/packages/autogen-studio/autogenstudio/cli.py +++ b/python/packages/autogen-studio/autogenstudio/cli.py @@ -131,6 +131,35 @@ def version(): typer.echo(f"AutoGen Studio CLI version: {VERSION}") +@app.command() +def lite( + team: Optional[str] = None, + host: str = "127.0.0.1", + port: int = 8080, + auto_open: bool = True, + session_name: str = "Lite Session", +): + """ + Launch AutoGen Studio in lightweight mode for quick experimentation. + + Args: + team (str, optional): Path to team JSON/YAML file. If not provided, uses a default team. + host (str): Host to run on. Defaults to 127.0.0.1. + port (int): Port to run on. Defaults to 8080. + auto_open (bool): Auto-open browser. Defaults to True. + session_name (str): Name for the auto-created session. + """ + from autogenstudio.lite import LiteStudio + + # Create and start studio instance + studio = LiteStudio(team=team, host=host, port=port, auto_open=auto_open, session_name=session_name) + + try: + studio.start() # Blocking call for CLI + except KeyboardInterrupt: + studio.stop() + + def run(): app() diff --git a/python/packages/autogen-studio/autogenstudio/database/db_manager.py b/python/packages/autogen-studio/autogenstudio/database/db_manager.py index f42eb95c26e5..f3eceebbe891 100644 --- a/python/packages/autogen-studio/autogenstudio/database/db_manager.py +++ b/python/packages/autogen-studio/autogenstudio/database/db_manager.py @@ -8,16 +8,24 @@ from sqlalchemy import exc, inspect, text from sqlmodel import Session, SQLModel, and_, create_engine, select -from ..datamodel import Response, Team +from ..datamodel import BaseDBModel, Response, Team from ..teammanager import TeamManager from .schema_manager import SchemaManager class CustomJSONEncoder(json.JSONEncoder): - def default(self, obj): - if hasattr(obj, "get_secret_value") and callable(obj.get_secret_value): - return obj.get_secret_value() - return super().default(obj) + def default(self, o): + if hasattr(o, "get_secret_value") and callable(o.get_secret_value): + return o.get_secret_value() + # Handle datetime objects + if isinstance(o, datetime): + return o.isoformat() + # Handle Enum objects + import enum + + if isinstance(o, enum.Enum): + return o.value + return super().default(o) class DatabaseManager: @@ -79,11 +87,14 @@ def initialize_database(self, auto_upgrade: bool = False, force_init_alembic: bo return Response(message="Failed to initialize migrations", status=False) # Handle existing database - if auto_upgrade or self._should_auto_upgrade(): + if auto_upgrade: logger.info("Checking database schema...") if self.schema_manager.ensure_schema_up_to_date(): return Response(message="Database schema is up to date", status=True) return Response(message="Database upgrade failed", status=False) + elif self._should_auto_upgrade(): + logger.info("Schema changes detected, but auto_upgrade is disabled. Skipping migration check.") + logger.info("To enable automatic migrations, set auto_upgrade=True") return Response(message="Database is ready", status=True) @@ -94,7 +105,7 @@ def initialize_database(self, auto_upgrade: bool = False, force_init_alembic: bo finally: self._init_lock.release() - def reset_db(self, recreate_tables: bool = True): + def reset_db(self, recreate_tables: bool = True) -> Response: """ Reset the database by dropping all tables and optionally recreating them. @@ -151,7 +162,7 @@ def reset_db(self, recreate_tables: bool = True): self._init_lock.release() logger.info("Database reset lock released") - def upsert(self, model: SQLModel, return_json: bool = True) -> Response: + def upsert(self, model: BaseDBModel, return_json: bool = True) -> Response: """Create or update an entity Args: @@ -199,7 +210,7 @@ def _model_to_dict(self, model_obj): def get( self, - model_class: SQLModel, + model_class: type[BaseDBModel], filters: dict | None = None, return_json: bool = False, order: str = "desc", @@ -211,7 +222,7 @@ def get( status_message = "" try: - statement = select(model_class) + statement = select(model_class) # type: ignore if filters: conditions = [getattr(model_class, col) == value for col, value in filters.items()] statement = statement.where(and_(*conditions)) @@ -231,7 +242,7 @@ def get( return Response(message=status_message, status=status, data=result) - def delete(self, model_class: SQLModel, filters: dict = None) -> Response: + def delete(self, model_class: type[BaseDBModel], filters: dict | None = None) -> Response: """Delete an entity""" status_message = "" status = True @@ -239,8 +250,8 @@ def delete(self, model_class: SQLModel, filters: dict = None) -> Response: with Session(self.engine) as session: try: if "sqlite" in str(self.engine.url): - session.exec(text("PRAGMA foreign_keys=ON")) - statement = select(model_class) + session.exec(text("PRAGMA foreign_keys=ON")) # type: ignore + statement = select(model_class) # type: ignore if filters: conditions = [getattr(model_class, col) == value for col, value in filters.items()] statement = statement.where(and_(*conditions)) @@ -326,7 +337,7 @@ async def import_teams_from_directory( { "status": result.status, "message": result.message, - "id": result.data.get("id") if result.status else None, + "id": result.data.get("id") if result.data and result.data is not None else None, } ) @@ -342,7 +353,8 @@ async def import_teams_from_directory( async def _check_team_exists(self, config: dict, user_id: str) -> Optional[Team]: """Check if identical team config already exists""" - teams = self.get(Team, {"user_id": user_id}).data + response = self.get(Team, {"user_id": user_id}) + teams = response.data if response.status and response.data is not None else [] for team in teams: if team.component == config: diff --git a/python/packages/autogen-studio/autogenstudio/datamodel/__init__.py b/python/packages/autogen-studio/autogenstudio/datamodel/__init__.py index 594982515aa0..ff5f0a08dbc7 100644 --- a/python/packages/autogen-studio/autogenstudio/datamodel/__init__.py +++ b/python/packages/autogen-studio/autogenstudio/datamodel/__init__.py @@ -1,4 +1,4 @@ -from .db import Gallery, Message, Run, RunStatus, Session, Settings, Team +from .db import BaseDBModel, Gallery, Message, Run, RunStatus, Session, Settings, Team from .types import ( EnvironmentVariable, GalleryComponents, diff --git a/python/packages/autogen-studio/autogenstudio/datamodel/db.py b/python/packages/autogen-studio/autogenstudio/datamodel/db.py index a7a1e2dcf973..ea8f427ed5b5 100644 --- a/python/packages/autogen-studio/autogenstudio/datamodel/db.py +++ b/python/packages/autogen-studio/autogenstudio/datamodel/db.py @@ -2,13 +2,14 @@ from datetime import datetime from enum import Enum -from typing import List, Optional, Union +from typing import Any, Dict, List, Optional, Union from autogen_core import ComponentModel -from pydantic import ConfigDict, SecretStr -from sqlalchemy import ForeignKey, Integer, String +from pydantic import ConfigDict, SecretStr, field_validator +from sqlalchemy import ForeignKey, Integer from sqlmodel import JSON, Column, DateTime, Field, SQLModel, func +from .eval import EvalJudgeCriteria, EvalRunResult, EvalRunStatus, EvalScore, EvalTask from .types import ( GalleryComponents, GalleryConfig, @@ -20,35 +21,41 @@ ) -class Team(SQLModel, table=True): - __table_args__ = {"sqlite_autoincrement": True} +class BaseDBModel(SQLModel, table=False): + """ + Base model with common fields for all database tables. + Not a table itself - meant to be inherited by concrete model classes. + """ + + __abstract__ = True + + # Common fields present in all database tables id: Optional[int] = Field(default=None, primary_key=True) + created_at: datetime = Field( default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), server_default=func.now()), - ) # pylint: disable=not-callable + sa_type=DateTime(timezone=True), # type: ignore[assignment] + sa_column_kwargs={"server_default": func.now(), "nullable": True}, + ) + updated_at: datetime = Field( default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), onupdate=func.now()), - ) # pylint: disable=not-callable + sa_type=DateTime(timezone=True), # type: ignore[assignment] + sa_column_kwargs={"onupdate": func.now(), "nullable": True}, + ) + user_id: Optional[str] = None version: Optional[str] = "0.0.1" + + +class Team(BaseDBModel, table=True): + __table_args__ = {"sqlite_autoincrement": True} component: Union[ComponentModel, dict] = Field(sa_column=Column(JSON)) -class Message(SQLModel, table=True): +class Message(BaseDBModel, table=True): __table_args__ = {"sqlite_autoincrement": True} - id: Optional[int] = Field(default=None, primary_key=True) - created_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), server_default=func.now()), - ) # pylint: disable=not-callable - updated_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), onupdate=func.now()), - ) # pylint: disable=not-callable - user_id: Optional[str] = None - version: Optional[str] = "0.0.1" + config: Union[MessageConfig, dict] = Field( default_factory=lambda: MessageConfig(source="", content=""), sa_column=Column(JSON) ) @@ -60,22 +67,18 @@ class Message(SQLModel, table=True): message_meta: Optional[Union[MessageMeta, dict]] = Field(default={}, sa_column=Column(JSON)) -class Session(SQLModel, table=True): +class Session(BaseDBModel, table=True): __table_args__ = {"sqlite_autoincrement": True} - id: Optional[int] = Field(default=None, primary_key=True) - created_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), server_default=func.now()), - ) # pylint: disable=not-callable - updated_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), onupdate=func.now()), - ) # pylint: disable=not-callable - user_id: Optional[str] = None - version: Optional[str] = "0.0.1" team_id: Optional[int] = Field(default=None, sa_column=Column(Integer, ForeignKey("team.id", ondelete="CASCADE"))) name: Optional[str] = None + @field_validator("created_at", "updated_at", mode="before") + @classmethod + def parse_datetime(cls, value: Union[str, datetime]) -> datetime: + if isinstance(value, str): + return datetime.fromisoformat(value.replace("Z", "+00:00")) + return value + class RunStatus(str, Enum): CREATED = "created" @@ -85,21 +88,12 @@ class RunStatus(str, Enum): STOPPED = "stopped" -class Run(SQLModel, table=True): +class Run(BaseDBModel, table=True): """Represents a single execution run within a session""" __table_args__ = {"sqlite_autoincrement": True} - id: Optional[int] = Field(default=None, primary_key=True) - created_at: datetime = Field( - default_factory=datetime.now, sa_column=Column(DateTime(timezone=True), server_default=func.now()) - ) - updated_at: datetime = Field( - default_factory=datetime.now, sa_column=Column(DateTime(timezone=True), onupdate=func.now()) - ) - session_id: Optional[int] = Field( - default=None, sa_column=Column(Integer, ForeignKey("session.id", ondelete="CASCADE"), nullable=False) - ) + session_id: int = Field(sa_column=Column(Integer, ForeignKey("session.id", ondelete="CASCADE"), nullable=False)) status: RunStatus = Field(default=RunStatus.CREATED) # Store the original user task @@ -111,32 +105,20 @@ class Run(SQLModel, table=True): team_result: Union[TeamResult, dict] = Field(default=None, sa_column=Column(JSON)) error_message: Optional[str] = None - version: Optional[str] = "0.0.1" messages: Union[List[Message], List[dict]] = Field(default_factory=list, sa_column=Column(JSON)) model_config = ConfigDict(json_encoders={datetime: lambda v: v.isoformat()}) # type: ignore[call-arg] - user_id: Optional[str] = None -class Gallery(SQLModel, table=True): +class Gallery(BaseDBModel, table=True): __table_args__ = {"sqlite_autoincrement": True} - id: Optional[int] = Field(default=None, primary_key=True) - created_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), server_default=func.now()), - ) # pylint: disable=not-callable - updated_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), onupdate=func.now()), - ) # pylint: disable=not-callable - user_id: Optional[str] = None - version: Optional[str] = "0.0.1" + config: Union[GalleryConfig, dict] = Field( default_factory=lambda: GalleryConfig( id="", name="", metadata=GalleryMetadata(author="", version=""), - components=GalleryComponents(agents=[], models=[], tools=[], terminations=[], teams=[]), + components=GalleryComponents(agents=[], models=[], tools=[], terminations=[], teams=[], workbenches=[]), ), sa_column=Column(JSON), ) @@ -149,17 +131,64 @@ class Gallery(SQLModel, table=True): ) # type: ignore[call-arg] -class Settings(SQLModel, table=True): +class Settings(BaseDBModel, table=True): __table_args__ = {"sqlite_autoincrement": True} - id: Optional[int] = Field(default=None, primary_key=True) - created_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), server_default=func.now()), - ) # pylint: disable=not-callable - updated_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), onupdate=func.now()), - ) # pylint: disable=not-callable - user_id: Optional[str] = None - version: Optional[str] = "0.0.1" + config: Union[SettingsConfig, dict] = Field(default_factory=SettingsConfig, sa_column=Column(JSON)) + + +# --- Evaluation system database models --- + + +class EvalTaskDB(BaseDBModel, table=True): + """Database model for storing evaluation tasks.""" + + __table_args__ = {"sqlite_autoincrement": True} + + name: str = "Unnamed Task" + description: str = "" + config: Union[EvalTask, dict] = Field(sa_column=Column(JSON)) + + +class EvalCriteriaDB(BaseDBModel, table=True): + """Database model for storing evaluation criteria.""" + + __table_args__ = {"sqlite_autoincrement": True} + + name: str = "Unnamed Criteria" + description: str = "" + config: Union[EvalJudgeCriteria, dict] = Field(sa_column=Column(JSON)) + + +class EvalRunDB(BaseDBModel, table=True): + """Database model for tracking evaluation runs.""" + + __table_args__ = {"sqlite_autoincrement": True} + + name: str = "Unnamed Evaluation Run" + description: str = "" + + # References to related components + task_id: Optional[int] = Field( + default=None, sa_column=Column(Integer, ForeignKey("evaltaskdb.id", ondelete="SET NULL")) + ) + + # Serialized configurations for runner and judge + runner_config: Union[ComponentModel, dict] = Field(sa_column=Column(JSON)) + judge_config: Union[ComponentModel, dict] = Field(sa_column=Column(JSON)) + + # List of criteria IDs or embedded criteria configs + criteria_configs: List[Union[EvalJudgeCriteria, dict]] = Field(default_factory=list, sa_column=Column(JSON)) + + # Run status and timing information + status: EvalRunStatus = Field(default=EvalRunStatus.PENDING) + start_time: Optional[datetime] = Field(default=None) + end_time: Optional[datetime] = Field(default=None) + + # Results (updated as they become available) + run_result: Union[EvalRunResult, dict] = Field(default=None, sa_column=Column(JSON)) + + score_result: Union[EvalScore, dict] = Field(default=None, sa_column=Column(JSON)) + + # Additional metadata + error_message: Optional[str] = None diff --git a/python/packages/autogen-studio/autogenstudio/datamodel/eval.py b/python/packages/autogen-studio/autogenstudio/datamodel/eval.py new file mode 100644 index 000000000000..4a1c3ec7db06 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/datamodel/eval.py @@ -0,0 +1,82 @@ +# datamodel/eval.py +from datetime import datetime +from enum import Enum +from typing import Any, Dict, List, Optional, Sequence +from uuid import UUID, uuid4 + +from autogen_agentchat.base import TaskResult +from autogen_core import Image +from pydantic import BaseModel +from sqlmodel import Field + + +class EvalTask(BaseModel): + """Definition of a task to be evaluated.""" + + task_id: UUID | str = Field(default_factory=uuid4) + input: str | Sequence[str | Image] + name: str = "" + description: str = "" + expected_outputs: Optional[List[Any]] = None + metadata: Dict[str, Any] = {} + + +class EvalRunResult(BaseModel): + """Result of an evaluation run.""" + + result: TaskResult | None = None + status: bool = False + start_time: Optional[datetime] = Field(default=datetime.now()) + end_time: Optional[datetime] = None + error: Optional[str] = None + + +class EvalDimensionScore(BaseModel): + """Score for a single evaluation dimension.""" + + dimension: str + score: float + reason: str + max_value: float + min_value: float + + +class EvalScore(BaseModel): + """Composite score from evaluation.""" + + overall_score: Optional[float] = None + dimension_scores: List[EvalDimensionScore] = [] + reason: Optional[str] = None + max_value: float = 10.0 + min_value: float = 0.0 + metadata: Dict[str, Any] = {} + + +class EvalJudgeCriteria(BaseModel): + """Criteria for judging evaluation results.""" + + dimension: str + prompt: str + max_value: float = 10.0 + min_value: float = 0.0 + metadata: Dict[str, Any] = {} + + +class EvalRunStatus(str, Enum): + """Status of an evaluation run.""" + + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELED = "canceled" + + +class EvalResult(BaseModel): + """Result of an evaluation run.""" + + task_id: UUID | str + # runner_id: UUID | str + status: EvalRunStatus = EvalRunStatus.PENDING + start_time: Optional[datetime] = Field(default=datetime.now()) + end_time: Optional[datetime] = None diff --git a/python/packages/autogen-studio/autogenstudio/datamodel/types.py b/python/packages/autogen-studio/autogenstudio/datamodel/types.py index 71a62a7124e8..cab12555caf4 100644 --- a/python/packages/autogen-studio/autogenstudio/datamodel/types.py +++ b/python/packages/autogen-studio/autogenstudio/datamodel/types.py @@ -1,17 +1,18 @@ # from dataclasses import Field from datetime import datetime -from typing import Any, Dict, List, Literal, Optional +from typing import Any, Dict, List, Literal, Optional, Sequence from autogen_agentchat.base import TaskResult -from autogen_agentchat.messages import BaseChatMessage +from autogen_agentchat.messages import ChatMessage, TextMessage from autogen_core import ComponentModel +from autogen_core.models import UserMessage from autogen_ext.models.openai import OpenAIChatCompletionClient -from pydantic import BaseModel, ConfigDict, SecretStr +from pydantic import BaseModel, ConfigDict, Field, SecretStr class MessageConfig(BaseModel): source: str - content: str + content: str | ChatMessage | Sequence[ChatMessage] | None message_type: Optional[str] = "text" @@ -21,9 +22,17 @@ class TeamResult(BaseModel): duration: float -class LLMCallEventMessage(BaseChatMessage): +class LLMCallEventMessage(TextMessage): source: str = "llm_call_event" - content: str + + def to_text(self) -> str: + return self.content + + def to_model_text(self) -> str: + return self.content + + def to_model_message(self) -> UserMessage: + raise NotImplementedError("This message type is not supported.") class MessageMeta(BaseModel): @@ -61,6 +70,7 @@ class GalleryComponents(BaseModel): tools: List[ComponentModel] terminations: List[ComponentModel] teams: List[ComponentModel] + workbenches: List[ComponentModel] class GalleryConfig(BaseModel): @@ -87,6 +97,9 @@ class UISettings(BaseModel): show_llm_call_events: bool = False expanded_messages_by_default: bool = True show_agent_flow_by_default: bool = True + human_input_timeout_minutes: int = Field( + default=3, ge=1, le=30, description="Human input timeout in minutes (1-30)" + ) class SettingsConfig(BaseModel): diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/__init__.py b/python/packages/autogen-studio/autogenstudio/eval/__init__.py similarity index 100% rename from python/packages/autogen-ext/src/autogen_ext/tools/__init__.py rename to python/packages/autogen-studio/autogenstudio/eval/__init__.py diff --git a/python/packages/autogen-studio/autogenstudio/eval/judges.py b/python/packages/autogen-studio/autogenstudio/eval/judges.py new file mode 100644 index 000000000000..e98b800a80df --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/eval/judges.py @@ -0,0 +1,267 @@ +import asyncio +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional, Tuple + +from autogen_core import CancellationToken, Component, ComponentBase +from autogen_core.models import ChatCompletionClient, UserMessage +from loguru import logger +from pydantic import BaseModel +from typing_extensions import Self + +from ..datamodel.eval import EvalDimensionScore, EvalJudgeCriteria, EvalRunResult, EvalScore, EvalTask + + +class BaseEvalJudgeConfig(BaseModel): + """Base configuration for evaluation judges.""" + + name: str = "Base Judge" + description: str = "" + metadata: Dict[str, Any] = {} + + +class BaseEvalJudge(ABC, ComponentBase[BaseEvalJudgeConfig]): + """Abstract base class for evaluation judges.""" + + component_type = "eval_judge" + + def __init__(self, name: str = "Base Judge", description: str = "", metadata: Optional[Dict[str, Any]] = None): + self.name = name + self.description = description + self.metadata = metadata or {} + + @abstractmethod + async def judge( + self, + task: EvalTask, + result: EvalRunResult, + criteria: List[EvalJudgeCriteria], + cancellation_token: Optional[CancellationToken] = None, + ) -> EvalScore: + """Judge the result of an evaluation run.""" + pass + + def _to_config(self) -> BaseEvalJudgeConfig: + """Convert the judge configuration to a configuration object for serialization.""" + return BaseEvalJudgeConfig(name=self.name, description=self.description, metadata=self.metadata) + + +class LLMEvalJudgeConfig(BaseEvalJudgeConfig): + """Configuration for LLMEvalJudge.""" + + model_client: Any # ComponentModel + + +class LLMEvalJudge(BaseEvalJudge, Component[LLMEvalJudgeConfig]): + """Judge that uses an LLM to evaluate results.""" + + component_config_schema = LLMEvalJudgeConfig + component_type = "eval_judge" + component_provider_override = "autogenstudio.eval.judges.LLMEvalJudge" + + def __init__( + self, + model_client: ChatCompletionClient, + name: str = "LLM Judge", + description: str = "Evaluates results using an LLM", + metadata: Optional[Dict[str, Any]] = None, + ): + super().__init__(name, description, metadata) + self.model_client = model_client + + async def judge( + self, + task: EvalTask, + result: EvalRunResult, + criteria: List[EvalJudgeCriteria], + cancellation_token: Optional[CancellationToken] = None, + ) -> EvalScore: + """Judge the result using an LLM.""" + # Create a score object + score = EvalScore(max_value=10.0) + + # Judge each dimension in parallel + dimension_score_tasks = [] + for criterion in criteria: + dimension_score_tasks.append(self._judge_dimension(task, result, criterion, cancellation_token)) + + dimension_scores = await asyncio.gather(*dimension_score_tasks) + score.dimension_scores = dimension_scores + + # Calculate overall score (average of dimension scores) + valid_scores = [ds.score for ds in dimension_scores if ds.score is not None] + if valid_scores: + score.overall_score = sum(valid_scores) / len(valid_scores) + + return score + + async def _judge_dimension( + self, + task: EvalTask, + result: EvalRunResult, + criterion: EvalJudgeCriteria, + cancellation_token: Optional[CancellationToken] = None, + ) -> EvalDimensionScore: + """Judge a specific dimension.""" + # Format task and result for the LLM + task_description = self._format_task(task) + result_description = result.model_dump() + + # Create the prompt + prompt = f""" + You are evaluating the quality of a system response to a task. + Task: {task_description}Response: {result_description} + Evaluation criteria: {criterion.dimension} + {criterion.prompt} + Score the response on a scale from {criterion.min_value} to {criterion.max_value}. + First, provide a detailed explanation of your evaluation. + Then, give your final score as a single number between 0 and {criterion.max_value}. + Format your answer should be a json for the EvalDimensionScore class: + {{ + "dimension": "{criterion.dimension}", + "reason": "", + "score": + }} + Please ensure the score is a number between {criterion.min_value} and {criterion.max_value}. + If you cannot evaluate the response, please return a score of null. + If the response is not relevant, please return a score of 0. + If the response is perfect, please return a score of {criterion.max_value}. + If the response is not relevant, please return a score of 0. + If the response is perfect, please return a score of {criterion.max_value}. + """ + + # Get judgment from LLM + model_input = [] + text_message = UserMessage(content=prompt, source="user") + model_input.append(text_message) + + # Run with the model client in the same format as used in runners + model_result = await self.model_client.create( + messages=model_input, + cancellation_token=cancellation_token, + json_output=EvalDimensionScore, + ) + + # Extract content from the response + model_response = model_result.content if isinstance(model_result.content, str) else str(model_result.content) + + try: + # validate response string as EvalDimensionScore + model_response = EvalDimensionScore.model_validate_json(model_response) + return model_response + except Exception as e: + logger.warning(f"Failed to parse LLM response: {e}", model_result.content) + return EvalDimensionScore( + dimension=criterion.dimension, + reason="Failed to parse response", + score=0.0, + max_value=criterion.max_value, + min_value=criterion.min_value, + ) + + def _format_task(self, task: EvalTask) -> str: + """Format the task for the LLM.""" + task_parts = [] + + if task.description: + task_parts.append(task.description) + if isinstance(task.input, str): + task_parts.append(task.input) + elif isinstance(task.input, list): + task_parts.append("\n".join(str(x) for x in task.input if isinstance(x, str))) + + return "\n".join(task_parts) + + def _parse_judgment(self, judgment_text: str, max_value: float) -> Tuple[str, Optional[float]]: + """Parse judgment text to extract explanation and score.""" + explanation = "" + score = None + + # Simple parsing - could be improved with regex + lines = judgment_text.split("\n") + for line in lines: + if line.strip().lower().startswith("explanation:"): + explanation = line.split(":", 1)[1].strip() + elif line.strip().lower().startswith("score:"): + try: + score_str = line.split(":", 1)[1].strip() + score = float(score_str) + # Ensure score is within bounds + score = min(max(score, 0), max_value) + except (ValueError, IndexError): + pass + + return explanation, score + + def _to_config(self) -> LLMEvalJudgeConfig: + """Convert to configuration object including model client configuration.""" + base_config = super()._to_config() + return LLMEvalJudgeConfig( + name=base_config.name, + description=base_config.description, + metadata=base_config.metadata, + model_client=self.model_client.dump_component(), + ) + + @classmethod + def _from_config(cls, config: LLMEvalJudgeConfig) -> Self: + """Create from configuration object with serialized model client.""" + model_client = ChatCompletionClient.load_component(config.model_client) + return cls( + model_client=model_client, name=config.name, description=config.description, metadata=config.metadata + ) + + +# # Usage example +# async def example_usage(): +# # Create a model client +# from autogen_ext.models import OpenAIChatCompletionClient + +# model_client = OpenAIChatCompletionClient( +# model="gpt-4", +# api_key="your-api-key" +# ) + +# # Create a judge +# llm_judge = LLMEvalJudge(model_client=model_client) + +# # Serialize the judge to a ComponentModel +# judge_config = llm_judge.dump_component() +# print(f"Serialized judge: {judge_config}") + +# # Deserialize back to a LLMEvalJudge +# deserialized_judge = LLMEvalJudge.load_component(judge_config) + +# # Create criteria for evaluation +# criteria = [ +# EvalJudgeCriteria( +# dimension="relevance", +# prompt="Evaluate how relevant the response is to the query.", +# min_value=0, +# max_value=10 +# ), +# EvalJudgeCriteria( +# dimension="accuracy", +# prompt="Evaluate the factual accuracy of the response.", +# min_value=0, +# max_value=10 +# ) +# ] + +# # Create a mock task and result +# task = EvalTask( +# id="task-123", +# name="Sample Task", +# description="A sample task for evaluation", +# input="What is the capital of France?" +# ) + +# result = EvalRunResult( +# status=True, +# result={ +# "messages": [{"content": "The capital of France is Paris.", "source": "model"}] +# } +# ) + +# # Run the evaluation +# score = await deserialized_judge.judge(task, result, criteria) +# print(f"Evaluation score: {score}") diff --git a/python/packages/autogen-studio/autogenstudio/eval/orchestrator.py b/python/packages/autogen-studio/autogenstudio/eval/orchestrator.py new file mode 100644 index 000000000000..ec5c711241f0 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/eval/orchestrator.py @@ -0,0 +1,789 @@ +import asyncio +import uuid +from datetime import datetime +from pdb import run +from typing import Any, Dict, List, Optional, TypedDict, Union + +from loguru import logger +from pydantic import BaseModel + +from ..database.db_manager import DatabaseManager +from ..datamodel.db import EvalCriteriaDB, EvalRunDB, EvalTaskDB +from ..datamodel.eval import EvalJudgeCriteria, EvalRunResult, EvalRunStatus, EvalScore, EvalTask +from .judges import BaseEvalJudge +from .runners import BaseEvalRunner + + +class DimensionScore(TypedDict): + score: Optional[float] + reason: Optional[str] + + +class RunEntry(TypedDict): + id: str + name: str + task_name: str + runner_type: str + overall_score: Optional[float] + scores: List[Optional[float]] + reasons: Optional[List[Optional[str]]] + + +class TabulatedResults(TypedDict): + dimensions: List[str] + runs: List[RunEntry] + + +class EvalOrchestrator: + """ + Orchestrator for evaluation runs. + + This class manages the lifecycle of evaluation tasks, criteria, and runs. + It can operate with or without a database manager for persistence. + """ + + def __init__(self, db_manager: Optional[DatabaseManager] = None): + """ + Initialize the orchestrator. + + Args: + db_manager: Optional database manager for persistence. + If None, data is stored in memory only. + """ + self._db_manager = db_manager + + # In-memory storage (used when db_manager is None) + self._tasks: Dict[str, EvalTask] = {} + self._criteria: Dict[str, EvalJudgeCriteria] = {} + self._runs: Dict[str, Dict[str, Any]] = {} + + # Active runs tracking + self._active_runs: Dict[str, asyncio.Task] = {} + + # ----- Task Management ----- + + async def create_task(self, task: EvalTask) -> str: + """ + Create a new evaluation task. + + Args: + task: The evaluation task to create + + Returns: + Task ID + """ + if not task.task_id: + task.task_id = str(uuid.uuid4()) + + if self._db_manager: + # Store in database + task_db = EvalTaskDB(name=task.name, description=task.description, config=task) + response = self._db_manager.upsert(task_db) + if not response.status: + logger.error(f"Failed to store task: {response.message}") + raise RuntimeError(f"Failed to store task: {response.message}") + task_id = str(response.data.get("id")) if response.data else str(task.task_id) + else: + # Store in memory + task_id = str(task.task_id) + self._tasks[task_id] = task + + return task_id + + async def get_task(self, task_id: str) -> Optional[EvalTask]: + """ + Retrieve an evaluation task by ID. + + Args: + task_id: The ID of the task to retrieve + + Returns: + The task if found, None otherwise + """ + if self._db_manager: + # Retrieve from database + response = self._db_manager.get(EvalTaskDB, filters={"id": int(task_id) if task_id.isdigit() else task_id}) + + if response.status and response.data and len(response.data) > 0: + task_data = response.data[0] + return ( + task_data.get("config") + if isinstance(task_data.get("config"), EvalTask) + else EvalTask.model_validate(task_data.get("config")) + ) + else: + # Retrieve from memory + return self._tasks.get(task_id) + + return None + + async def list_tasks(self) -> List[EvalTask]: + """ + List all available evaluation tasks. + + Returns: + List of evaluation tasks + """ + if self._db_manager: + # Retrieve from database + response = self._db_manager.get(EvalTaskDB) + + tasks = [] + if response.status and response.data: + for task_data in response.data: + config = task_data.get("config") + if config: + if isinstance(config, EvalTask): + tasks.append(config) + else: + tasks.append(EvalTask.model_validate(config)) + return tasks + else: + # Retrieve from memory + return list(self._tasks.values()) + + # ----- Criteria Management ----- + + async def create_criteria(self, criteria: EvalJudgeCriteria) -> str: + """ + Create new evaluation criteria. + + Args: + criteria: The evaluation criteria to create + + Returns: + Criteria ID + """ + criteria_id = str(uuid.uuid4()) + + if self._db_manager: + # Store in database + criteria_db = EvalCriteriaDB(name=criteria.dimension, description=criteria.prompt, config=criteria) + response = self._db_manager.upsert(criteria_db) + if not response.status: + logger.error(f"Failed to store criteria: {response.message}") + raise RuntimeError(f"Failed to store criteria: {response.message}") + criteria_id = str(response.data.get("id")) if response.data else criteria_id + else: + # Store in memory + self._criteria[criteria_id] = criteria + + return criteria_id + + async def get_criteria(self, criteria_id: str) -> Optional[EvalJudgeCriteria]: + """ + Retrieve evaluation criteria by ID. + + Args: + criteria_id: The ID of the criteria to retrieve + + Returns: + The criteria if found, None otherwise + """ + if self._db_manager: + # Retrieve from database + response = self._db_manager.get( + EvalCriteriaDB, filters={"id": int(criteria_id) if criteria_id.isdigit() else criteria_id} + ) + + if response.status and response.data and len(response.data) > 0: + criteria_data = response.data[0] + return ( + criteria_data.get("config") + if isinstance(criteria_data.get("config"), EvalJudgeCriteria) + else EvalJudgeCriteria.model_validate(criteria_data.get("config")) + ) + else: + # Retrieve from memory + return self._criteria.get(criteria_id) + + return None + + async def list_criteria(self) -> List[EvalJudgeCriteria]: + """ + List all available evaluation criteria. + + Returns: + List of evaluation criteria + """ + if self._db_manager: + # Retrieve from database + response = self._db_manager.get(EvalCriteriaDB) + + criteria_list = [] + if response.status and response.data: + for criteria_data in response.data: + config = criteria_data.get("config") + if config: + if isinstance(config, EvalJudgeCriteria): + criteria_list.append(config) + else: + criteria_list.append(EvalJudgeCriteria.model_validate(config)) + return criteria_list + else: + # Retrieve from memory + return list(self._criteria.values()) + + # ----- Run Management ----- + + async def create_run( + self, + task: Union[str, EvalTask], + runner: BaseEvalRunner, + judge: BaseEvalJudge, + criteria: List[Union[str, EvalJudgeCriteria]], + name: str = "", + description: str = "", + ) -> str: + """ + Create a new evaluation run configuration. + + Args: + task: The task to evaluate (ID or task object) + runner: The runner to use for evaluation + judge: The judge to use for evaluation + criteria: List of criteria to use for evaluation (IDs or criteria objects) + name: Name for the run + description: Description for the run + + Returns: + Run ID + """ + # Resolve task + task_obj = None + if isinstance(task, str): + task_obj = await self.get_task(task) + if not task_obj: + raise ValueError(f"Task not found: {task}") + else: + task_obj = task + + # Resolve criteria + criteria_objs = [] + for criterion in criteria: + if isinstance(criterion, str): + criterion_obj = await self.get_criteria(criterion) + if not criterion_obj: + raise ValueError(f"Criteria not found: {criterion}") + criteria_objs.append(criterion_obj) + else: + criteria_objs.append(criterion) + + # Generate run ID + run_id = str(uuid.uuid4()) + + # Create run configuration + runner_config = runner.dump_component() if hasattr(runner, "dump_component") else runner._to_config() + judge_config = judge.dump_component() if hasattr(judge, "dump_component") else judge._to_config() + + if self._db_manager: + # Store in database + run_db = EvalRunDB( + name=name or f"Run {run_id}", + description=description, + task_id=int(task) if isinstance(task, str) and task.isdigit() else None, + runner_config=runner_config.model_dump(), + judge_config=judge_config.model_dump(), + criteria_configs=criteria_objs, + status=EvalRunStatus.PENDING, + ) + response = self._db_manager.upsert(run_db) + if not response.status: + logger.error(f"Failed to store run: {response.message}") + raise RuntimeError(f"Failed to store run: {response.message}") + run_id = str(response.data.get("id")) if response.data else run_id + else: + # Store in memory + self._runs[run_id] = { + "task": task_obj, + "runner_config": runner_config, + "judge_config": judge_config, + "criteria_configs": [c.model_dump() for c in criteria_objs], + "status": EvalRunStatus.PENDING, + "created_at": datetime.now(), + "run_result": None, + "score_result": None, + "name": name or f"Run {run_id}", + "description": description, + } + + return run_id + + async def start_run(self, run_id: str) -> None: + """ + Start an evaluation run. + + Args: + run_id: The ID of the run to start + """ + # Check if run is already active + if run_id in self._active_runs: + logger.warning(f"Run {run_id} is already active") + return + + # Start the run asynchronously + run_task = asyncio.create_task(self._execute_run(run_id)) + self._active_runs[run_id] = run_task + + # Update run status + await self._update_run_status(run_id, EvalRunStatus.RUNNING) + + async def _execute_run(self, run_id: str) -> None: + """ + Execute an evaluation run. + + Args: + run_id: The ID of the run to execute + """ + try: + # Get run configuration + run_config = await self._get_run_config(run_id) + if not run_config: + raise ValueError(f"Run not found: {run_id}") + + # Get task + task = run_config.get("task") + if not task: + raise ValueError(f"Task not found for run: {run_id}") + + # Initialize runner + runner_config = run_config.get("runner_config") + runner = BaseEvalRunner.load_component(runner_config) if runner_config else None + + # Initialize judge + judge_config = run_config.get("judge_config") + judge = BaseEvalJudge.load_component(judge_config) if judge_config else None + + if not runner or not judge: + raise ValueError(f"Runner or judge not found for run: {run_id}") + + # Initialize criteria + criteria_configs = run_config.get("criteria_configs") + criteria = [] + if criteria_configs: + criteria = [ + EvalJudgeCriteria.model_validate(c) if not isinstance(c, EvalJudgeCriteria) else c + for c in criteria_configs + ] + + # Execute runner + logger.info(f"Starting runner for run {run_id}") + start_time = datetime.now() + run_result = await runner.run(task) + + # Update run result + await self._update_run_result(run_id, run_result) + + if not run_result.status: + logger.error(f"Runner failed for run {run_id}: {run_result.error}") + await self._update_run_status(run_id, EvalRunStatus.FAILED) + return + + # Execute judge + logger.info(f"Starting judge for run {run_id}") + score_result = await judge.judge(task, run_result, criteria) + + # Update score result + await self._update_score_result(run_id, score_result) + + # Update run status + end_time = datetime.now() + await self._update_run_completed(run_id, start_time, end_time) + + logger.info(f"Run {run_id} completed successfully") + + except Exception as e: + logger.exception(f"Error executing run {run_id}: {str(e)}") + await self._update_run_error(run_id, str(e)) + finally: + # Remove from active runs + if run_id in self._active_runs: + del self._active_runs[run_id] + + async def get_run_status(self, run_id: str) -> Optional[EvalRunStatus]: + """ + Get the status of an evaluation run. + + Args: + run_id: The ID of the run + + Returns: + The run status if found, None otherwise + """ + run_config = await self._get_run_config(run_id) + return run_config.get("status") if run_config else None + + async def get_run_result(self, run_id: str) -> Optional[EvalRunResult]: + """ + Get the result of an evaluation run. + + Args: + run_id: The ID of the run + + Returns: + The run result if found, None otherwise + """ + run_config = await self._get_run_config(run_id) + if not run_config: + return None + + run_result = run_config.get("run_result") + if not run_result: + return None + + return run_result if isinstance(run_result, EvalRunResult) else EvalRunResult.model_validate(run_result) + + async def get_run_score(self, run_id: str) -> Optional[EvalScore]: + """ + Get the score of an evaluation run. + + Args: + run_id: The ID of the run + + Returns: + The run score if found, None otherwise + """ + run_config = await self._get_run_config(run_id) + if not run_config: + return None + + score_result = run_config.get("score_result") + if not score_result: + return None + + return score_result if isinstance(score_result, EvalScore) else EvalScore.model_validate(score_result) + + async def list_runs(self) -> List[Dict[str, Any]]: + """ + List all available evaluation runs. + + Returns: + List of run configurations + """ + if self._db_manager: + # Retrieve from database + response = self._db_manager.get(EvalRunDB) + + runs = [] + if response.status and response.data: + for run_data in response.data: + runs.append( + { + "id": run_data.get("id"), + "name": run_data.get("name"), + "status": run_data.get("status"), + "created_at": run_data.get("created_at"), + "updated_at": run_data.get("updated_at"), + } + ) + return runs + else: + # Retrieve from memory + return [ + { + "id": run_id, + "name": run_config.get("name"), + "status": run_config.get("status"), + "created_at": run_config.get("created_at"), + "updated_at": run_config.get("updated_at", run_config.get("created_at")), + } + for run_id, run_config in self._runs.items() + ] + + async def cancel_run(self, run_id: str) -> bool: + """ + Cancel an active evaluation run. + + Args: + run_id: The ID of the run to cancel + + Returns: + True if the run was cancelled, False otherwise + """ + # Check if run is active + if run_id not in self._active_runs: + logger.warning(f"Run {run_id} is not active") + return False + + # Cancel the run task + try: + self._active_runs[run_id].cancel() + await self._update_run_status(run_id, EvalRunStatus.CANCELED) + del self._active_runs[run_id] + return True + except Exception as e: + logger.error(f"Failed to cancel run {run_id}: {str(e)}") + return False + + # ----- Helper Methods ----- + + async def _get_run_config(self, run_id: str) -> Optional[Dict[str, Any]]: + """ + Get the configuration of an evaluation run. + + Args: + run_id: The ID of the run + + Returns: + The run configuration if found, None otherwise + """ + if self._db_manager: + # Retrieve from database + response = self._db_manager.get(EvalRunDB, filters={"id": int(run_id) if run_id.isdigit() else run_id}) + + if response.status and response.data and len(response.data) > 0: + run_data = response.data[0] + + # Get task + task = None + if run_data.get("task_id"): + task_response = self._db_manager.get(EvalTaskDB, filters={"id": run_data.get("task_id")}) + if task_response.status and task_response.data and len(task_response.data) > 0: + task_data = task_response.data[0] + task = ( + task_data.get("config") + if isinstance(task_data.get("config"), EvalTask) + else EvalTask.model_validate(task_data.get("config")) + ) + + return { + "task": task, + "runner_config": run_data.get("runner_config"), + "judge_config": run_data.get("judge_config"), + "criteria_configs": run_data.get("criteria_configs"), + "status": run_data.get("status"), + "run_result": run_data.get("run_result"), + "score_result": run_data.get("score_result"), + "name": run_data.get("name"), + "description": run_data.get("description"), + "created_at": run_data.get("created_at"), + "updated_at": run_data.get("updated_at"), + } + else: + # Retrieve from memory + return self._runs.get(run_id) + + return None + + async def _update_run_status(self, run_id: str, status: EvalRunStatus) -> None: + """ + Update the status of an evaluation run. + + Args: + run_id: The ID of the run + status: The new status + """ + if self._db_manager: + # Update in database + response = self._db_manager.get(EvalRunDB, filters={"id": int(run_id) if run_id.isdigit() else run_id}) + + if response.status and response.data and len(response.data) > 0: + run_data = response.data[0] + run_db = EvalRunDB.model_validate(run_data) + run_db.status = status + run_db.updated_at = datetime.now() + self._db_manager.upsert(run_db) + else: + # Update in memory + if run_id in self._runs: + self._runs[run_id]["status"] = status + self._runs[run_id]["updated_at"] = datetime.now() + + async def _update_run_result(self, run_id: str, run_result: EvalRunResult) -> None: + """ + Update the result of an evaluation run. + + Args: + run_id: The ID of the run + run_result: The run result + """ + if self._db_manager: + # Update in database + response = self._db_manager.get(EvalRunDB, filters={"id": int(run_id) if run_id.isdigit() else run_id}) + + if response.status and response.data and len(response.data) > 0: + run_data = response.data[0] + run_db = EvalRunDB.model_validate(run_data) + run_db.run_result = run_result + run_db.updated_at = datetime.now() + self._db_manager.upsert(run_db) + else: + # Update in memory + if run_id in self._runs: + self._runs[run_id]["run_result"] = run_result + self._runs[run_id]["updated_at"] = datetime.now() + + async def _update_score_result(self, run_id: str, score_result: EvalScore) -> None: + """ + Update the score of an evaluation run. + + Args: + run_id: The ID of the run + score_result: The score result + """ + if self._db_manager: + # Update in database + response = self._db_manager.get(EvalRunDB, filters={"id": int(run_id) if run_id.isdigit() else run_id}) + + if response.status and response.data and len(response.data) > 0: + run_data = response.data[0] + run_db = EvalRunDB.model_validate(run_data) + run_db.score_result = score_result + run_db.updated_at = datetime.now() + self._db_manager.upsert(run_db) + else: + # Update in memory + if run_id in self._runs: + self._runs[run_id]["score_result"] = score_result + self._runs[run_id]["updated_at"] = datetime.now() + + async def _update_run_completed(self, run_id: str, start_time: datetime, end_time: datetime) -> None: + """ + Update a run as completed. + + Args: + run_id: The ID of the run + start_time: The start time + end_time: The end time + """ + if self._db_manager: + # Update in database + response = self._db_manager.get(EvalRunDB, filters={"id": int(run_id) if run_id.isdigit() else run_id}) + + if response.status and response.data and len(response.data) > 0: + run_data = response.data[0] + run_db = EvalRunDB.model_validate(run_data) + run_db.status = EvalRunStatus.COMPLETED + run_db.start_time = start_time + run_db.end_time = end_time + run_db.updated_at = datetime.now() + self._db_manager.upsert(run_db) + else: + # Update in memory + if run_id in self._runs: + self._runs[run_id]["status"] = EvalRunStatus.COMPLETED + self._runs[run_id]["start_time"] = start_time + self._runs[run_id]["end_time"] = end_time + self._runs[run_id]["updated_at"] = datetime.now() + + async def _update_run_error(self, run_id: str, error_message: str) -> None: + """ + Update a run with an error. + + Args: + run_id: The ID of the run + error_message: The error message + """ + if self._db_manager: + # Update in database + response = self._db_manager.get(EvalRunDB, filters={"id": int(run_id) if run_id.isdigit() else run_id}) + + if response.status and response.data and len(response.data) > 0: + run_data = response.data[0] + run_db = EvalRunDB.model_validate(run_data) + run_db.status = EvalRunStatus.FAILED + run_db.error_message = error_message + run_db.end_time = datetime.now() + run_db.updated_at = datetime.now() + self._db_manager.upsert(run_db) + else: + # Update in memory + if run_id in self._runs: + self._runs[run_id]["status"] = EvalRunStatus.FAILED + self._runs[run_id]["error_message"] = error_message + self._runs[run_id]["end_time"] = datetime.now() + self._runs[run_id]["updated_at"] = datetime.now() + + async def tabulate_results(self, run_ids: List[str], include_reasons: bool = False) -> TabulatedResults: + """ + Generate a tabular representation of evaluation results across runs. + + This method collects scores across different runs and organizes them by + dimension, making it easy to create visualizations like radar charts. + + Args: + run_ids: List of run IDs to include in the tabulation + include_reasons: Whether to include scoring reasons in the output + + Returns: + A dictionary with structured data suitable for visualization + """ + result: TabulatedResults = {"dimensions": [], "runs": []} + + # Parallelize fetching of run configs and scores + fetch_tasks = [] + for run_id in run_ids: + fetch_tasks.append(self._get_run_config(run_id)) + fetch_tasks.append(self.get_run_score(run_id)) + + # Wait for all fetches to complete + fetch_results = await asyncio.gather(*fetch_tasks) + + # Process fetched data + dimensions_set = set() + run_data = {} + + for i in range(0, len(fetch_results), 2): + run_id = run_ids[i // 2] + run_config = fetch_results[i] + score = fetch_results[i + 1] + + # Store run data for later processing + run_data[run_id] = (run_config, score) + + # Collect dimensions + if score and score.dimension_scores: + for dim_score in score.dimension_scores: + dimensions_set.add(dim_score.dimension) + + # Convert dimensions to sorted list + result["dimensions"] = sorted(list(dimensions_set)) + + # Process each run's data + for run_id, (run_config, score) in run_data.items(): + if not run_config or not score: + continue + + # Determine runner type + runner_type = "unknown" + if run_config.get("runner_config"): + runner_config = run_config.get("runner_config") + if runner_config is not None and "provider" in runner_config: + if "ModelEvalRunner" in runner_config["provider"]: + runner_type = "model" + elif "TeamEvalRunner" in runner_config["provider"]: + runner_type = "team" + + # Get task name + task = run_config.get("task") + task_name = task.name if task else "Unknown Task" + + # Create run entry + run_entry: RunEntry = { + "id": run_id, + "name": run_config.get("name", f"Run {run_id}"), + "task_name": task_name, + "runner_type": runner_type, + "overall_score": score.overall_score, + "scores": [], + "reasons": [] if include_reasons else None, + } + + # Build dimension lookup map for O(1) access + dim_map = {ds.dimension: ds for ds in score.dimension_scores} + + # Populate scores aligned with dimensions + for dim in result["dimensions"]: + dim_score = dim_map.get(dim) + if dim_score: + run_entry["scores"].append(dim_score.score) + if include_reasons: + run_entry["reasons"].append(dim_score.reason) # type: ignore + else: + run_entry["scores"].append(None) + if include_reasons: + run_entry["reasons"].append(None) # type: ignore + + result["runs"].append(run_entry) + + return result diff --git a/python/packages/autogen-studio/autogenstudio/eval/runners.py b/python/packages/autogen-studio/autogenstudio/eval/runners.py new file mode 100644 index 000000000000..bbb1f4813c4d --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/eval/runners.py @@ -0,0 +1,201 @@ +from abc import ABC, abstractmethod +from datetime import datetime +from typing import Any, Dict, Optional, Sequence, Type, Union + +from autogen_agentchat.base import TaskResult, Team +from autogen_agentchat.messages import ChatMessage, MultiModalMessage, TextMessage +from autogen_core import CancellationToken, Component, ComponentBase, ComponentModel, Image +from autogen_core.models import ChatCompletionClient, UserMessage +from pydantic import BaseModel +from typing_extensions import Self + +from ..datamodel.eval import EvalRunResult, EvalTask + + +class BaseEvalRunnerConfig(BaseModel): + """Base configuration for evaluation runners.""" + + name: str + description: str = "" + metadata: Dict[str, Any] = {} + + +class BaseEvalRunner(ABC, ComponentBase[BaseEvalRunnerConfig]): + """Base class for evaluation runners that defines the interface for running evaluations. + + This class provides the core interface that all evaluation runners must implement. + Subclasses should implement the run method to define how a specific evaluation is executed. + """ + + component_type = "eval_runner" + + def __init__(self, name: str, description: str = "", metadata: Optional[Dict[str, Any]] = None): + self.name = name + self.description = description + self.metadata = metadata or {} + + @abstractmethod + async def run(self, task: EvalTask, cancellation_token: Optional[CancellationToken] = None) -> EvalRunResult: + """Run the evaluation on the provided task and return a result. + + Args: + task: The task to evaluate + cancellation_token: Optional token to cancel the evaluation + + Returns: + EvaluationResult: The result of the evaluation + """ + pass + + def _to_config(self) -> BaseEvalRunnerConfig: + """Convert the runner configuration to a configuration object for serialization.""" + return BaseEvalRunnerConfig(name=self.name, description=self.description, metadata=self.metadata) + + +class ModelEvalRunnerConfig(BaseEvalRunnerConfig): + """Configuration for ModelEvalRunner.""" + + model_client: ComponentModel + + +class ModelEvalRunner(BaseEvalRunner, Component[ModelEvalRunnerConfig]): + """Evaluation runner that uses a single LLM to process tasks. + + This runner sends the task directly to a model client and returns the response. + """ + + component_config_schema = ModelEvalRunnerConfig + component_type = "eval_runner" + component_provider_override = "autogenstudio.eval.runners.ModelEvalRunner" + + def __init__( + self, + model_client: ChatCompletionClient, + name: str = "Model Runner", + description: str = "Evaluates tasks using a single LLM", + metadata: Optional[Dict[str, Any]] = None, + ): + super().__init__(name, description, metadata) + self.model_client = model_client + + async def run(self, task: EvalTask, cancellation_token: Optional[CancellationToken] = None) -> EvalRunResult: + """Run the task with the model client and return the result.""" + # Create initial result object + result = EvalRunResult() + + try: + model_input = [] + if isinstance(task.input, str): + text_message = UserMessage(content=task.input, source="user") + model_input.append(text_message) + elif isinstance(task.input, list): + message_content = [x for x in task.input] + model_input.append(UserMessage(content=message_content, source="user")) + # Run with the model + model_result = await self.model_client.create(messages=model_input, cancellation_token=cancellation_token) + + model_response = model_result.content if isinstance(model_result, str) else model_result.model_dump() + + task_result = TaskResult( + messages=[TextMessage(content=str(model_response), source="model")], + ) + result = EvalRunResult(result=task_result, status=True, start_time=datetime.now(), end_time=datetime.now()) + + except Exception as e: + result = EvalRunResult(status=False, error=str(e), end_time=datetime.now()) + + return result + + def _to_config(self) -> ModelEvalRunnerConfig: + """Convert to configuration object including model client configuration.""" + base_config = super()._to_config() + return ModelEvalRunnerConfig( + name=base_config.name, + description=base_config.description, + metadata=base_config.metadata, + model_client=self.model_client.dump_component(), + ) + + @classmethod + def _from_config(cls, config: ModelEvalRunnerConfig) -> Self: + """Create from configuration object with serialized model client.""" + model_client = ChatCompletionClient.load_component(config.model_client) + return cls( + name=config.name, + description=config.description, + metadata=config.metadata, + model_client=model_client, + ) + + +class TeamEvalRunnerConfig(BaseEvalRunnerConfig): + """Configuration for TeamEvalRunner.""" + + team: ComponentModel + + +class TeamEvalRunner(BaseEvalRunner, Component[TeamEvalRunnerConfig]): + """Evaluation runner that uses a team of agents to process tasks. + + This runner creates and runs a team based on a team configuration. + """ + + component_config_schema = TeamEvalRunnerConfig + component_type = "eval_runner" + component_provider_override = "autogenstudio.eval.runners.TeamEvalRunner" + + def __init__( + self, + team: Union[Team, ComponentModel], + name: str = "Team Runner", + description: str = "Evaluates tasks using a team of agents", + metadata: Optional[Dict[str, Any]] = None, + ): + super().__init__(name, description, metadata) + self._team = team if isinstance(team, Team) else Team.load_component(team) + + async def run(self, task: EvalTask, cancellation_token: Optional[CancellationToken] = None) -> EvalRunResult: + """Run the task with the team and return the result.""" + # Create initial result object + result = EvalRunResult() + + try: + team_task: Sequence[ChatMessage] = [] + if isinstance(task.input, str): + team_task.append(TextMessage(content=task.input, source="user")) + if isinstance(task.input, list): + for message in task.input: + if isinstance(message, str): + team_task.append(TextMessage(content=message, source="user")) + elif isinstance(message, Image): + team_task.append(MultiModalMessage(source="user", content=[message])) + + # Run task with team + team_result = await self._team.run(task=team_task, cancellation_token=cancellation_token) + + result = EvalRunResult(result=team_result, status=True, start_time=datetime.now(), end_time=datetime.now()) + + except Exception as e: + result = EvalRunResult(status=False, error=str(e), end_time=datetime.now()) + + return result + + def _to_config(self) -> TeamEvalRunnerConfig: + """Convert to configuration object including team configuration.""" + base_config = super()._to_config() + return TeamEvalRunnerConfig( + name=base_config.name, + description=base_config.description, + metadata=base_config.metadata, + team=self._team.dump_component(), + ) + + @classmethod + def _from_config(cls, config: TeamEvalRunnerConfig) -> Self: + """Create from configuration object with serialized team configuration.""" + return cls( + team=Team.load_component(config.team), + name=config.name, + description=config.description, + metadata=config.metadata, + ) diff --git a/python/packages/autogen-studio/autogenstudio/gallery/builder.py b/python/packages/autogen-studio/autogenstudio/gallery/builder.py index 6d5067b205d6..55a124367dd4 100644 --- a/python/packages/autogen-studio/autogenstudio/gallery/builder.py +++ b/python/packages/autogen-studio/autogenstudio/gallery/builder.py @@ -1,19 +1,29 @@ import os -from datetime import datetime +import tempfile from typing import List, Optional -import anthropic from autogen_agentchat.agents import AssistantAgent, UserProxyAgent -from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination -from autogen_agentchat.teams import RoundRobinGroupChat, SelectorGroupChat +from autogen_agentchat.conditions import ( + HandoffTermination, + MaxMessageTermination, + SourceMatchTermination, + StopMessageTermination, + TextMentionTermination, + TextMessageTermination, + TimeoutTermination, + TokenUsageTermination, +) +from autogen_agentchat.teams import RoundRobinGroupChat, SelectorGroupChat, Swarm from autogen_core import ComponentModel from autogen_core.models import ModelInfo +from autogen_core.tools import StaticWorkbench from autogen_ext.agents.web_surfer import MultimodalWebSurfer from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor from autogen_ext.models.anthropic import AnthropicChatCompletionClient from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.models.openai._openai_client import AzureOpenAIChatCompletionClient from autogen_ext.tools.code_execution import PythonCodeExecutionTool +from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams, StreamableHttpServerParams from autogenstudio.datamodel import GalleryComponents, GalleryConfig, GalleryMetadata @@ -32,6 +42,7 @@ def __init__(self, id: str, name: str, url: Optional[str] = None): self.models: List[ComponentModel] = [] self.tools: List[ComponentModel] = [] self.terminations: List[ComponentModel] = [] + self.workbenches: List[ComponentModel] = [] # Default metadata self.metadata = GalleryMetadata( @@ -112,6 +123,13 @@ def add_termination( self.terminations.append(self._update_component_metadata(termination, label, description)) return self + def add_workbench( + self, workbench: ComponentModel, label: Optional[str] = None, description: Optional[str] = None + ) -> "GalleryBuilder": + """Add a workbench component to the gallery with optional custom label and description.""" + self.workbenches.append(self._update_component_metadata(workbench, label, description)) + return self + def build(self) -> GalleryConfig: """Build and return the complete gallery.""" # Update timestamps @@ -128,6 +146,7 @@ def build(self) -> GalleryConfig: models=self.models, tools=self.tools, terminations=self.terminations, + workbenches=self.workbenches, ), ) @@ -220,6 +239,64 @@ def create_default_gallery() -> GalleryConfig: description="Termination condition that ends the conversation when either a message contains 'TERMINATE' or the maximum number of messages is reached.", ) + # Add examples of new termination conditions + + # StopMessageTermination - terminates when a StopMessage is received + stop_msg_term = StopMessageTermination() + builder.add_termination( + stop_msg_term.dump_component(), + label="Stop Message Termination", + description="Terminates the conversation when a StopMessage is received from any agent.", + ) + + # TokenUsageTermination - terminates based on token usage limits + token_usage_term = TokenUsageTermination(max_total_token=1000, max_prompt_token=800, max_completion_token=200) + builder.add_termination( + token_usage_term.dump_component(), + label="Token Usage Termination", + description="Terminates the conversation when token usage limits are reached (1000 total, 800 prompt, 200 completion).", + ) + + # TimeoutTermination - terminates after a specified duration + timeout_term = TimeoutTermination(timeout_seconds=300) # 5 minutes + builder.add_termination( + timeout_term.dump_component(), + label="Timeout Termination", + description="Terminates the conversation after 5 minutes (300 seconds) have elapsed.", + ) + + # HandoffTermination - terminates when handoff to specific target occurs + handoff_term = HandoffTermination(target="user_proxy") + builder.add_termination( + handoff_term.dump_component(), + label="Handoff Termination", + description="Terminates the conversation when a handoff to 'user_proxy' is detected.", + ) + + # SourceMatchTermination - terminates when specific sources respond + source_match_term = SourceMatchTermination(sources=["assistant_agent", "critic_agent"]) + builder.add_termination( + source_match_term.dump_component(), + label="Source Match Termination", + description="Terminates the conversation when either 'assistant_agent' or 'critic_agent' responds.", + ) + + # TextMessageTermination - terminates on TextMessage from specific source + text_msg_term = TextMessageTermination(source="assistant_agent") + builder.add_termination( + text_msg_term.dump_component(), + label="Text Message Termination", + description="Terminates the conversation when a TextMessage is received from 'assistant_agent'.", + ) + + # Create a complex termination combining multiple conditions + complex_term = (token_usage_term | timeout_term) & (calc_text_term | stop_msg_term) + builder.add_termination( + complex_term.dump_component(), + label="Complex Termination", + description="Complex termination: (token usage OR timeout) AND (text mention 'TERMINATE' OR stop message).", + ) + # Create calculator team calc_team = RoundRobinGroupChat(participants=[calc_assistant], termination_condition=calc_or_term) builder.add_team( @@ -243,6 +320,31 @@ def create_default_gallery() -> GalleryConfig: description="A team with 2 agents - an AssistantAgent (with a calculator tool) and a CriticAgent in a SelectorGroupChat team.", ) + # Create Swarm team - agents with handoff capabilities + # Alice agent with handoff to Bob + alice_agent = AssistantAgent( + name="Alice", + system_message="You are Alice, a helpful assistant. You specialize in general questions. If someone asks about technical topics or needs detailed analysis, hand off to Bob by saying 'Let me hand this over to Bob for a detailed analysis.'", + model_client=base_model, + handoffs=["Bob"], + ) + + # Bob agent with handoff back to Alice + bob_agent = AssistantAgent( + name="Bob", + system_message="You are Bob, a technical specialist. You handle detailed technical analysis. If the conversation becomes general or the user needs basic assistance, hand off to Alice by saying 'Let me hand this back to Alice for general assistance.'", + model_client=base_model, + handoffs=["Alice"], + ) + + # Create simple Swarm team with handoff-based conversation + swarm_team = Swarm(participants=[alice_agent, bob_agent], termination_condition=calc_or_term) + builder.add_team( + swarm_team.dump_component(), + label="Swarm Team", + description="A team with 2 agents (Alice and Bob) that use handoff messages to transfer conversation control between agents based on expertise.", + ) + # Create web surfer agent websurfer_agent = MultimodalWebSurfer( name="websurfer_agent", @@ -413,9 +515,116 @@ def create_default_gallery() -> GalleryConfig: description="A team with 3 agents - a Research Assistant that performs web searches and analyzes information, a Verifier that ensures research quality and completeness, and a Summary Agent that provides a detailed markdown summary of the research as a report to the user.", ) + # Add workbenches to the gallery + + # Create a static workbench with basic tools + static_workbench = StaticWorkbench(tools=[tools.calculator_tool, tools.fetch_webpage_tool]) + builder.add_workbench( + static_workbench.dump_component(), + label="Basic Tools Workbench", + description="A static workbench containing basic tools like calculator and webpage fetcher for common tasks.", + ) + + # Create an MCP workbench for fetching web content using mcp-server-fetch + # Note: This requires uv to be installed (comes with uv package manager) + fetch_server_params = StdioServerParams( + command="uv", + args=["tool", "run", "mcp-server-fetch"], + read_timeout_seconds=60, + ) + mcp_workbench = McpWorkbench(server_params=fetch_server_params) + builder.add_workbench( + mcp_workbench.dump_component(), + label="MCP Fetch Workbench", + description="An MCP workbench that provides web content fetching capabilities using the mcp-server-fetch MCP server. Allows agents to fetch and read content from web pages and APIs.", + ) + + # Create an MCP workbench with StreamableHttpServerParams for HTTP-based MCP servers + # Note: This is an example - adjust URL and authentication as needed + streamable_server_params = StreamableHttpServerParams( + url="http://localhost:8005/mcp", + headers={"Authorization": "Bearer your-api-key", "Content-Type": "application/json"}, + timeout=30, + sse_read_timeout=60 * 5, + terminate_on_close=True, + ) + streamable_mcp_workbench = McpWorkbench(server_params=streamable_server_params) + builder.add_workbench( + streamable_mcp_workbench.dump_component(), + label="MCP Streamable HTTP Workbench", + description="An MCP workbench that connects to HTTP-based MCP servers using Server-Sent Events (SSE). Suitable for cloud-hosted MCP services and custom HTTP MCP implementations.", + ) + + # Create an MCP workbench for filesystem operations + # Note: This requires npx to be installed and allows access to specified directories + + # Use cross-platform paths for filesystem access + user_home = os.path.expanduser("~") + temp_dir = tempfile.gettempdir() # Cross-platform temp directory + + filesystem_server_params = StdioServerParams( + command="npx", + args=["-y", "@modelcontextprotocol/server-filesystem", user_home, temp_dir], + read_timeout_seconds=60, + ) + filesystem_mcp_workbench = McpWorkbench(server_params=filesystem_server_params) + builder.add_workbench( + filesystem_mcp_workbench.dump_component(), + label="MCP Filesystem Workbench", + description="An MCP workbench that provides filesystem access capabilities using the @modelcontextprotocol/server-filesystem MCP server. Allows agents to read, write, and manage files and directories within specified allowed paths.", + ) + + # Create an MCP workbench for testing with everything server + # Note: This requires npx to be installed and provides comprehensive MCP testing tools + everything_server_params = StdioServerParams( + command="npx", + args=["-y", "@modelcontextprotocol/server-everything"], + read_timeout_seconds=60, + ) + everything_mcp_workbench = McpWorkbench(server_params=everything_server_params) + builder.add_workbench( + everything_mcp_workbench.dump_component(), + label="MCP Test Server", + description="An MCP workbench that provides comprehensive testing tools using the @modelcontextprotocol/server-everything MCP server. Includes various tools for testing MCP functionality, protocol features, and capabilities.", + ) + return builder.build() +def create_default_lite_team(): + """Create a simple default team for lite mode - a basic assistant with calculator tool.""" + import json + import os + import tempfile + + # model clients require API keys to be set in the environment or passed in + # as arguments. For testing purposes, we set them to "test" if not already set. + for key in ["OPENAI_API_KEY", "AZURE_OPENAI_API_KEY", "ANTHROPIC_API_KEY"]: + if not os.environ.get(key): + os.environ[key] = "test" + + # Create base model client + base_model = OpenAIChatCompletionClient(model="gpt-4o-mini") + + # Create assistant agent with calculator tool + assistant = AssistantAgent( + name="assistant", + model_client=base_model, + tools=[tools.calculator_tool], + ) + + # Create termination condition + termination = TextMentionTermination(text="TERMINATE") | MaxMessageTermination(max_messages=5) + + # Create simple round robin team + team = RoundRobinGroupChat(participants=[assistant], termination_condition=termination) + + # Create temporary file with team data + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + json.dump(team.dump_component().model_dump(), f, indent=2) + return f.name + + if __name__ == "__main__": # Create and save the gallery gallery = create_default_gallery() diff --git a/python/packages/autogen-studio/autogenstudio/lite/__init__.py b/python/packages/autogen-studio/autogenstudio/lite/__init__.py new file mode 100644 index 000000000000..9f7c560b52a3 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/lite/__init__.py @@ -0,0 +1,3 @@ +from .studio import LiteStudio + +__all__ = ["LiteStudio"] diff --git a/python/packages/autogen-studio/autogenstudio/lite/studio.py b/python/packages/autogen-studio/autogenstudio/lite/studio.py new file mode 100644 index 000000000000..94b25cd85b6e --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/lite/studio.py @@ -0,0 +1,245 @@ +import json +import os +import subprocess +import tempfile +import threading +import time +import webbrowser +from pathlib import Path +from typing import Any, Dict, Union + +import uvicorn +from autogen_core import ComponentModel + + +class LiteStudio: + """ + Core class for managing AutoGen Studio lite mode instances. + Supports both file-based and programmatic team configurations. + """ + + def __init__( + self, + team: Union[str, Path, Dict[str, Any], ComponentModel, None] = None, + host: str = "127.0.0.1", + port: int = 8080, + session_name: str = "Lite Session", + auto_open: bool = True, + ): + """ + Initialize LiteStudio instance. + + Args: + team: Team configuration - can be: + - str: Path to team JSON file + - Path: Path object to team JSON file + - Dict[str, Any]: Team configuration dictionary + - ComponentModel: AutoGen ComponentModel instance + - None: Creates default team + host: Host to run server on + port: Port to run server on + session_name: Name for the auto-created session + auto_open: Whether to auto-open browser + """ + self.host = host + self.port = port + self.session_name = session_name + self.auto_open = auto_open + self.server_process = None + self.server_thread = None # Handle team loading + self.team_file_path = self._load_team(team) + + def _load_team(self, team: Union[str, Path, Dict[str, Any], ComponentModel, None]) -> str: + """ + Load team from file path, object, or create default. + Returns the file path to the team JSON. + Args: + team: Can be file path (str/Path), dict, ComponentModel, or None + """ + if team is None: + # Create default team + from autogenstudio.gallery.builder import create_default_lite_team + + return create_default_lite_team() + + elif isinstance(team, (str, Path)): + # File path provided + team_path = Path(team) + if not team_path.exists(): + raise FileNotFoundError(f"Team file not found: {team_path}") + return str(team_path.absolute()) + + elif isinstance(team, dict): + # Team dict provided - save to temp file + temp_file = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) + try: + json.dump(team, temp_file, indent=2) + temp_file.flush() + return temp_file.name + finally: + temp_file.close() + + elif isinstance(team, ComponentModel): + # ComponentModel - use model_dump directly + team_dict = team.model_dump() + temp_file = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) + try: + json.dump(team_dict, temp_file, indent=2) + temp_file.flush() + return temp_file.name + finally: + temp_file.close() + + else: + # Try to serialize other team objects + team_dict = None + + # Try dump_component() method (AutoGen teams) + if hasattr(team, "dump_component"): + component = team.dump_component() + if hasattr(component, "model_dump"): + team_dict = component.model_dump() + elif hasattr(component, "dict"): + team_dict = component.dict() + else: + team_dict = dict(component) + + # Try model_dump() method (Pydantic v2) + elif hasattr(team, "model_dump"): + team_dict = team.model_dump() + + # Try dict() method (Pydantic v1) + elif hasattr(team, "dict"): + team_dict = team.dict() + + if team_dict is None: + raise ValueError( + f"Cannot serialize team object of type {type(team)}. " + f"Expected: file path, dict, ComponentModel, or object with dump_component()/model_dump()/dict() method." + ) + + # Save serialized team to temp file + temp_file = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) + try: + json.dump(team_dict, temp_file, indent=2) + temp_file.flush() + return temp_file.name + finally: + temp_file.close() + + def _get_env_file_path(self) -> str: + """Get path for environment variables file.""" + app_dir = os.path.join(os.path.expanduser("~"), ".autogenstudio") + if not os.path.exists(app_dir): + os.makedirs(app_dir, exist_ok=True) + return os.path.join(app_dir, "temp_env_vars.env") + + def _setup_environment(self) -> str: + """ + Setup environment variables for lite mode. + Returns path to the environment file. + """ + env_vars = { + "AUTOGENSTUDIO_HOST": self.host, + "AUTOGENSTUDIO_PORT": str(self.port), + "AUTOGENSTUDIO_LITE_MODE": "true", + "AUTOGENSTUDIO_API_DOCS": "false", + "AUTOGENSTUDIO_AUTH_DISABLED": "true", + "AUTOGENSTUDIO_LITE_SESSION_NAME": self.session_name, + "AUTOGENSTUDIO_LITE_TEAM_FILE": self.team_file_path, + "AUTOGENSTUDIO_DATABASE_URI": "sqlite:///:memory:", + } + + env_file_path = self._get_env_file_path() + with open(env_file_path, "w") as temp_env: + for key, value in env_vars.items(): + temp_env.write(f"{key}={value}\n") + + return env_file_path + + def _setup_browser_opening(self): + """Setup browser opening in a separate thread.""" + if self.auto_open: + + def open_browser(): + time.sleep(3) # Wait for server startup + url = f"http://{self.host}:{self.port}/lite" + webbrowser.open(url) + + threading.Thread(target=open_browser, daemon=True).start() + + def start(self, background: bool = False): + """ + Start the lite studio server. + Args: + background: If True, run server in background thread + """ + # Check if already running + if self.server_thread and self.server_thread.is_alive(): + raise RuntimeError("LiteStudio is already running") + + # Setup environment + env_file_path = self._setup_environment() + + # Setup browser opening + self._setup_browser_opening() + + if background: + # Run server in background thread + def run_server(): + uvicorn.run( + "autogenstudio.web.app:app", + host=self.host, + port=self.port, + workers=1, + env_file=env_file_path, + ) + + self.server_thread = threading.Thread(target=run_server, daemon=True) + self.server_thread.start() + else: + # Run server in foreground (blocking) + uvicorn.run( + "autogenstudio.web.app:app", + host=self.host, + port=self.port, + workers=1, + env_file=env_file_path, + ) + + def stop(self): + """Stop the lite studio server.""" + if self.server_thread and self.server_thread.is_alive(): + # For background threads, we can't easily stop uvicorn + # This is a limitation - in production you'd want proper shutdown + self.server_thread.join(timeout=5) + self.server_thread = None + + def __enter__(self): + """Context manager entry - start in background.""" + self.start(background=True) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): # type: ignore + """Context manager exit - stop server.""" + self.stop() + + @classmethod + def shutdown_port(cls, port: int): + """ + Utility to shutdown any process running on the specified port. + Args: + port: Port number to shutdown + """ + try: + # Try to find and kill process on port (Unix/Linux/Mac) + result = subprocess.run(["lsof", "-ti", f":{port}"], capture_output=True, text=True) + + if result.returncode == 0 and result.stdout.strip(): + pids = result.stdout.strip().split("\n") + for pid in pids: + subprocess.run(["kill", "-9", pid], check=False) + + except (subprocess.SubprocessError, FileNotFoundError): + # lsof might not be available on all systems + pass diff --git a/python/packages/autogen-studio/autogenstudio/mcp/callbacks.py b/python/packages/autogen-studio/autogenstudio/mcp/callbacks.py new file mode 100644 index 000000000000..417f0f846c66 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/mcp/callbacks.py @@ -0,0 +1,186 @@ +import asyncio +import uuid +from datetime import datetime, timezone +from typing import Any, Dict, Tuple + +from loguru import logger +from mcp.shared.context import RequestContext +from mcp.shared.session import RequestResponder +from mcp.types import ( + ClientResult, + CreateMessageRequestParams, + CreateMessageResult, + ElicitRequestParams, + ElicitResult, + ErrorData, + ServerNotification, + ServerRequest, + TextContent, +) + +from .utils import extract_real_error, serialize_for_json +from .wsbridge import MCPWebSocketBridge + + +def create_message_handler(bridge: MCPWebSocketBridge): + """Create a message handler callback that streams MCP protocol messages to the UI""" + + async def message_handler( + message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, + ) -> None: + try: + if isinstance(message, Exception): + await bridge.on_mcp_activity( + "error", f"Protocol error: {str(message)}", {"details": extract_real_error(message)} + ) + + elif hasattr(message, "method"): + method = getattr(message, "method", "unknown") + params = getattr(message, "params", None) + + await bridge.on_mcp_activity( + "protocol", + f"MCP {method}", + { + "method": method, + "params": serialize_for_json(params) if params else None, + "message_type": type(message).__name__, + }, + ) + + else: + await bridge.on_mcp_activity( + "protocol", + f"{type(message).__name__}", + { + "message_type": type(message).__name__, + "content": serialize_for_json(message) if hasattr(message, "model_dump") else str(message), + }, + ) + + except Exception as e: + logger.error(f"Error in message handler: {extract_real_error(e)}") + + return message_handler + + +def create_sampling_callback(bridge: MCPWebSocketBridge): + """Create a sampling callback that handles AI sampling requests from tools""" + + async def sampling_callback( + context: RequestContext[Any, Any, Any], + params: CreateMessageRequestParams, + ) -> CreateMessageResult | ErrorData: + try: + request_id = str(uuid.uuid4()) + + await bridge.on_mcp_activity( + "sampling", + f"Tool requested AI sampling for {len(params.messages)} message(s)", + { + "request_id": request_id, + "params": serialize_for_json(params.model_dump()), + "context": "Tool is requesting AI to generate a response", + }, + ) + + dummy_response = CreateMessageResult( + role="assistant", + content=TextContent( + type="text", + text="[AutoGen Studio Default Sampling Response - This is a placeholder response for AI sampling requests. In a production setup, this would be handled by your configured LLM.]", + ), + model="autogen-studio-default", + ) + + await bridge.on_mcp_activity( + "sampling", + "Provided default sampling response to tool", + { + "request_id": request_id, + "response": serialize_for_json(dummy_response.model_dump()), + "note": "This is a placeholder response - configure an LLM for real sampling", + }, + ) + + logger.info("Handled sampling request with default response") + return dummy_response + + except Exception as e: + error_msg = extract_real_error(e) + logger.error(f"Error in sampling callback: {error_msg}") + + await bridge.on_mcp_activity("error", f"Sampling callback error: {error_msg}", {"error": error_msg}) + + return ErrorData(code=-32603, message=f"Sampling failed: {error_msg}") + + return sampling_callback + + +def create_elicitation_callback( + bridge: MCPWebSocketBridge, +) -> Tuple[Any, Dict[str, asyncio.Future[ElicitResult | ErrorData]]]: + """Create an elicitation callback that handles user input requests from tools""" + + async def elicitation_callback( + context: RequestContext[Any, Any, Any], + params: ElicitRequestParams, + ) -> ElicitResult | ErrorData: + try: + request_id = str(uuid.uuid4()) + + await bridge.on_mcp_activity( + "elicitation", + f"Tool requesting user input: {params.message}", + { + "request_id": request_id, + "message": params.message, + "requestedSchema": serialize_for_json(params.requestedSchema) if params.requestedSchema else None, + "context": "Tool is requesting additional information from user", + }, + ) + + await bridge.on_elicitation_request( + request_id, + params.message, + serialize_for_json(params.requestedSchema) if params.requestedSchema else None, + ) + + response_future: asyncio.Future[ElicitResult | ErrorData] = asyncio.Future() + bridge.pending_elicitations[request_id] = response_future + + try: + user_response = await asyncio.wait_for(response_future, timeout=60.0) + + await bridge.on_mcp_activity( + "elicitation", + "User responded to elicitation request", + {"request_id": request_id, "response": serialize_for_json(user_response), "status": "completed"}, + ) + + return user_response + + except asyncio.TimeoutError: + logger.warning(f"User did not respond to elicitation request {request_id} within 60 seconds") + error_msg = "User did not respond to elicitation request within 60 seconds" + + await bridge.on_mcp_activity( + "error", + f"Elicitation timeout: {error_msg}", + {"request_id": request_id, "error": error_msg, "timeout": 60}, + ) + + return ErrorData(code=-32603, message=error_msg) + + finally: + bridge.pending_elicitations.pop(request_id, None) + + except Exception as e: + error_msg = extract_real_error(e) + logger.error(f"Error in elicitation callback: {error_msg}") + + await bridge.on_mcp_activity("error", f"Elicitation callback error: {error_msg}", {"error": error_msg}) + + return ErrorData(code=-32603, message=f"Elicitation failed: {error_msg}") + + return elicitation_callback, bridge.pending_elicitations diff --git a/python/packages/autogen-studio/autogenstudio/mcp/client.py b/python/packages/autogen-studio/autogenstudio/mcp/client.py new file mode 100644 index 000000000000..a8a076ee73fd --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/mcp/client.py @@ -0,0 +1,121 @@ +from typing import Any, Dict, Protocol + +from mcp import ClientSession + +from .utils import McpOperationError, extract_real_error, serialize_for_json + + +class MCPEventHandler(Protocol): + """Interface for handling MCP events""" + + async def on_initialized(self, session_id: str, capabilities: Any) -> None: + """Called when MCP session is initialized""" + ... + + async def on_operation_result(self, operation: str, data: Dict[str, Any]) -> None: + """Called when an MCP operation completes successfully""" + ... + + async def on_operation_error(self, operation: str, error: str) -> None: + """Called when an MCP operation fails""" + ... + + async def on_mcp_activity(self, activity_type: str, message: str, details: Dict[str, Any]) -> None: + """Called for MCP protocol activity""" + ... + + async def on_elicitation_request(self, request_id: str, message: str, requested_schema: Any) -> None: + """Called when MCP requests user input""" + ... + + +class MCPClient: + """Handles MCP protocol operations independently of transport""" + + def __init__(self, session: ClientSession, session_id: str, event_handler: MCPEventHandler): + self.session = session + self.session_id = session_id + self.event_handler = event_handler + self._initialized = False + self._capabilities = None + + async def initialize(self) -> None: + """Initialize the MCP session""" + try: + initialize_result = await self.session.initialize() + + if initialize_result: + self._capabilities = initialize_result.capabilities + else: + self._capabilities = None + + self._initialized = True + + # Notify handler + await self.event_handler.on_initialized( + self.session_id, serialize_for_json(self._capabilities.model_dump()) if self._capabilities else None + ) + + except Exception as e: + await self.event_handler.on_operation_error("initialize", str(e)) + raise + + async def handle_operation(self, operation: Dict[str, Any]) -> None: + """Handle an MCP operation - this preserves the exact behavior of handle_mcp_operation""" + operation_type = operation.get("operation") + + try: + if operation_type == "list_tools": + result = await self.session.list_tools() + tools_data = [serialize_for_json(tool.model_dump()) for tool in result.tools] + await self.event_handler.on_operation_result("list_tools", {"tools": tools_data}) + + elif operation_type == "call_tool": + tool_name = operation.get("tool_name") + arguments = operation.get("arguments", {}) + if not tool_name: + raise McpOperationError("Tool name is required") + + result = await self.session.call_tool(tool_name, arguments) + await self.event_handler.on_operation_result( + "call_tool", {"tool_name": tool_name, "result": serialize_for_json(result.model_dump())} + ) + + elif operation_type == "list_resources": + result = await self.session.list_resources() + await self.event_handler.on_operation_result("list_resources", serialize_for_json(result.model_dump())) + + elif operation_type == "read_resource": + uri = operation.get("uri") + if not uri: + raise McpOperationError("Resource URI is required") + + result = await self.session.read_resource(uri) + await self.event_handler.on_operation_result("read_resource", serialize_for_json(result.model_dump())) + + elif operation_type == "list_prompts": + result = await self.session.list_prompts() + prompts_data = [serialize_for_json(prompt.model_dump()) for prompt in result.prompts] + await self.event_handler.on_operation_result("list_prompts", {"prompts": prompts_data}) + + elif operation_type == "get_prompt": + name = operation.get("name") + arguments = operation.get("arguments") + if not name: + raise McpOperationError("Prompt name is required") + + result = await self.session.get_prompt(name, arguments) + await self.event_handler.on_operation_result("get_prompt", serialize_for_json(result.model_dump())) + + else: + await self.event_handler.on_operation_error( + operation_type or "unknown", f"Unknown operation: {operation_type}" + ) + + except Exception as e: + real_error = extract_real_error(e) + await self.event_handler.on_operation_error(operation_type or "unknown", real_error) + + @property + def capabilities(self): + return self._capabilities diff --git a/python/packages/autogen-studio/autogenstudio/mcp/utils.py b/python/packages/autogen-studio/autogenstudio/mcp/utils.py new file mode 100644 index 000000000000..457957a5b7d0 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/mcp/utils.py @@ -0,0 +1,88 @@ +from typing import Any, List + +from fastapi import WebSocketDisconnect +from pydantic.networks import AnyUrl + + +class McpOperationError(Exception): + """Raised when MCP operation fails""" + + pass + + +def extract_real_error(e: Exception) -> str: + """Extract the real error message from potentially wrapped exceptions""" + error_parts: List[str] = [] + + # Handle ExceptionGroup (Python 3.11+) + if hasattr(e, "exceptions") and getattr(e, "exceptions", None): + exceptions_list = getattr(e, "exceptions", []) + for sub_exc in exceptions_list: + error_parts.append(f"{type(sub_exc).__name__}: {str(sub_exc)}") + + # Handle chained exceptions + elif hasattr(e, "__cause__") and e.__cause__: + current = e + while current: + error_parts.append(f"{type(current).__name__}: {str(current)}") + current = getattr(current, "__cause__", None) + + # Handle context exceptions + elif hasattr(e, "__context__") and e.__context__: + error_parts.append(f"Context: {type(e.__context__).__name__}: {str(e.__context__)}") + error_parts.append(f"Error: {type(e).__name__}: {str(e)}") + + # Default case + else: + error_parts.append(f"{type(e).__name__}: {str(e)}") + + return " | ".join(error_parts) + + +def serialize_for_json(obj: Any) -> Any: + """Convert objects to JSON-serializable format""" + if isinstance(obj, AnyUrl): + return str(obj) + elif isinstance(obj, dict): + return {str(k): serialize_for_json(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [serialize_for_json(item) for item in obj] + elif hasattr(obj, "model_dump"): + return serialize_for_json(obj.model_dump()) + else: + return obj + + +def is_websocket_disconnect(e: Exception) -> bool: + """Check if an exception (potentially nested) is a WebSocket disconnect""" + + def check_exception(exc: BaseException) -> bool: + if isinstance(exc, WebSocketDisconnect): + return True + + exc_name = type(exc).__name__ + exc_str = str(exc) + + if "WebSocketDisconnect" in exc_name or "NO_STATUS_RCVD" in exc_str: + return True + + # Recursively check ExceptionGroup + if hasattr(exc, "exceptions") and getattr(exc, "exceptions", None): + exceptions_list = getattr(exc, "exceptions", []) + for sub_exc in exceptions_list: + if check_exception(sub_exc): + return True + + # Check chained exceptions + if hasattr(exc, "__cause__") and exc.__cause__: + if check_exception(exc.__cause__): + return True + + # Check context exceptions + if hasattr(exc, "__context__") and exc.__context__: + if check_exception(exc.__context__): + return True + + return False + + return check_exception(e) diff --git a/python/packages/autogen-studio/autogenstudio/mcp/wsbridge.py b/python/packages/autogen-studio/autogenstudio/mcp/wsbridge.py new file mode 100644 index 000000000000..d723b085367e --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/mcp/wsbridge.py @@ -0,0 +1,215 @@ +import asyncio +import json +from datetime import datetime, timezone +from typing import Any, Dict, Optional + +from fastapi import WebSocket +from loguru import logger +from mcp import ClientSession +from mcp.types import ElicitResult, ErrorData + +from .client import MCPClient, MCPEventHandler +from .utils import extract_real_error, is_websocket_disconnect, serialize_for_json + + +class MCPWebSocketBridge(MCPEventHandler): + """Bridges WebSocket connections to MCP operations""" + + def __init__(self, websocket: WebSocket, session_id: str): + self.websocket = websocket + self.session_id = session_id + self.mcp_client: Optional[MCPClient] = None + self.pending_elicitations: Dict[str, asyncio.Future[ElicitResult | ErrorData]] = {} + self._running = True + + async def send_message(self, message: Dict[str, Any]) -> None: + """Send a message through the WebSocket""" + try: + from fastapi.websockets import WebSocketState + + if self.websocket.client_state == WebSocketState.CONNECTED: + serialized_message = serialize_for_json(message) + await self.websocket.send_json(serialized_message) + except Exception as e: + real_error = extract_real_error(e) + logger.error(f"Error sending WebSocket message: {real_error}") + + # Implement MCPEventHandler interface + async def on_initialized(self, session_id: str, capabilities: Any) -> None: + await self.send_message( + { + "type": "initialized", + "session_id": session_id, + "capabilities": capabilities, + "timestamp": datetime.now(timezone.utc).isoformat(), + } + ) + + async def on_operation_result(self, operation: str, data: Dict[str, Any]) -> None: + await self.send_message( + { + "type": "operation_result", + "operation": operation, + "data": data, + "timestamp": datetime.now(timezone.utc).isoformat(), + } + ) + + async def on_operation_error(self, operation: str, error: str) -> None: + await self.send_message( + { + "type": "operation_error", + "operation": operation, + "error": error, + "timestamp": datetime.now(timezone.utc).isoformat(), + } + ) + + async def on_mcp_activity(self, activity_type: str, message: str, details: Dict[str, Any]) -> None: + await self.send_message( + { + "type": "mcp_activity", + "activity_type": activity_type, + "message": message, + "details": details, + "session_id": self.session_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + } + ) + + async def on_elicitation_request(self, request_id: str, message: str, requested_schema: Any) -> None: + await self.send_message( + { + "type": "elicitation_request", + "request_id": request_id, + "message": message, + "requestedSchema": requested_schema, + "session_id": self.session_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + } + ) + + def set_mcp_client(self, mcp_client: MCPClient) -> None: + """Set the MCP client after initialization""" + self.mcp_client = mcp_client + + async def handle_websocket_message(self, message: Dict[str, Any]) -> None: + """Handle incoming WebSocket messages""" + message_type = message.get("type") + + # Update last activity for session tracking + from datetime import datetime, timezone + + from ..web.routes.mcp import active_sessions + + if self.session_id in active_sessions: + active_sessions[self.session_id]["last_activity"] = datetime.now(timezone.utc) + + if message_type == "operation": + # CRITICAL: Run in background task to avoid blocking message loop + # This preserves the exact behavior from the original code + if self.mcp_client: + asyncio.create_task(self.mcp_client.handle_operation(message)) + else: + await self.send_message( + { + "type": "error", + "error": "MCP client not initialized", + "timestamp": datetime.now(timezone.utc).isoformat(), + } + ) + + elif message_type == "ping": + await self.send_message({"type": "pong", "timestamp": datetime.now(timezone.utc).isoformat()}) + + elif message_type == "elicitation_response": + await self._handle_elicitation_response(message) + + else: + await self.send_message( + { + "type": "error", + "error": f"Unknown message type: {message_type}", + "timestamp": datetime.now(timezone.utc).isoformat(), + } + ) + + async def _handle_elicitation_response(self, message: Dict[str, Any]) -> None: + """Handle user response to elicitation request""" + request_id = message.get("request_id") + + if not request_id: + await self.send_message( + { + "type": "error", + "error": "Missing request_id in elicitation response", + "timestamp": datetime.now(timezone.utc).isoformat(), + } + ) + return + + if request_id in self.pending_elicitations: + try: + action = message.get("action", "cancel") + data = message.get("data", {}) + + if action == "accept": + result = ElicitResult(action="accept", content=data) + elif action == "decline": + result = ElicitResult(action="decline") + else: + result = ElicitResult(action="cancel") + + future = self.pending_elicitations[request_id] + if not future.done(): + future.set_result(result) + else: + logger.warning(f"Future for elicitation request {request_id} was already done") + + except Exception as e: + error_msg = extract_real_error(e) + logger.error(f"Error processing elicitation response: {error_msg}") + + future = self.pending_elicitations.get(request_id) + if future and not future.done(): + future.set_result( + ErrorData(code=-32603, message=f"Error processing elicitation response: {error_msg}") + ) + else: + logger.warning(f"Unknown elicitation request_id: {request_id}") + await self.send_message( + { + "type": "operation_error", + "error": f"Unknown elicitation request_id: {request_id}", + "timestamp": datetime.now(timezone.utc).isoformat(), + } + ) + + async def run(self) -> None: + """Main message loop""" + try: + while self._running: + try: + raw_message = await self.websocket.receive_text() + message = json.loads(raw_message) + await self.handle_websocket_message(message) + + except json.JSONDecodeError: + logger.warning(f"Invalid JSON received from session {self.session_id}") + await self.send_message( + { + "type": "error", + "error": "Invalid message format", + "timestamp": datetime.now(timezone.utc).isoformat(), + } + ) + + except Exception as e: + if not is_websocket_disconnect(e): + real_error = extract_real_error(e) + logger.error(f"Error in message loop: {real_error}") + raise + + def stop(self) -> None: + """Stop the bridge""" + self._running = False diff --git a/python/packages/autogen-studio/autogenstudio/teammanager/teammanager.py b/python/packages/autogen-studio/autogenstudio/teammanager/teammanager.py index 96dbaec1098c..f143a5d4db17 100644 --- a/python/packages/autogen-studio/autogenstudio/teammanager/teammanager.py +++ b/python/packages/autogen-studio/autogenstudio/teammanager/teammanager.py @@ -4,28 +4,33 @@ import os import time from pathlib import Path -from typing import AsyncGenerator, Callable, List, Optional, Union +from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, List, Optional, Sequence, Union import aiofiles import yaml from autogen_agentchat.agents import UserProxyAgent -from autogen_agentchat.base import TaskResult, Team -from autogen_agentchat.messages import AgentEvent, ChatMessage +from autogen_agentchat.base import TaskResult +from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage from autogen_agentchat.teams import BaseGroupChat -from autogen_core import EVENT_LOGGER_NAME, CancellationToken, Component, ComponentModel +from autogen_core import EVENT_LOGGER_NAME, CancellationToken, ComponentModel from autogen_core.logging import LLMCallEvent from ..datamodel.types import EnvironmentVariable, LLMCallEventMessage, TeamResult +from ..web.managers.run_context import RunContext logger = logging.getLogger(__name__) +SyncInputFunc = Callable[[str], str] +AsyncInputFunc = Callable[[str, Optional[CancellationToken]], Awaitable[str]] +InputFuncType = Union[SyncInputFunc, AsyncInputFunc] + class RunEventLogger(logging.Handler): """Event logger that queues LLMCallEvents for streaming""" def __init__(self): super().__init__() - self.events = asyncio.Queue() + self.events: asyncio.Queue[LLMCallEventMessage] = asyncio.Queue() def emit(self, record: logging.LogRecord): if isinstance(record.msg, LLMCallEvent): @@ -35,8 +40,12 @@ def emit(self, record: logging.LogRecord): class TeamManager: """Manages team operations including loading configs and running teams""" + def __init__(self): + self._team: Optional[BaseGroupChat] = None + self._run_context = RunContext() + @staticmethod - async def load_from_file(path: Union[str, Path]) -> dict: + async def load_from_file(path: Union[str, Path]) -> Any: """Load team configuration from JSON/YAML file""" path = Path(path) if not path.exists(): @@ -51,10 +60,10 @@ async def load_from_file(path: Union[str, Path]) -> dict: raise ValueError(f"Unsupported file format: {path.suffix}") @staticmethod - async def load_from_directory(directory: Union[str, Path]) -> List[dict]: + async def load_from_directory(directory: Union[str, Path]) -> List[Any]: """Load all team configurations from a directory""" directory = Path(directory) - configs = [] + configs: List[Any] = [] valid_extensions = {".json", ".yaml", ".yml"} for path in directory.iterdir(): @@ -69,8 +78,8 @@ async def load_from_directory(directory: Union[str, Path]) -> List[dict]: async def _create_team( self, - team_config: Union[str, Path, dict, ComponentModel], - input_func: Optional[Callable] = None, + team_config: Union[str, Path, Dict[str, Any], ComponentModel], + input_func: Optional[InputFuncType] = None, env_vars: Optional[List[EnvironmentVariable]] = None, ) -> BaseGroupChat: """Create team instance from config""" @@ -78,8 +87,10 @@ async def _create_team( config = await self.load_from_file(team_config) elif isinstance(team_config, dict): config = team_config - else: + elif isinstance(team_config, ComponentModel): config = team_config.model_dump() + else: + raise ValueError(f"Unsupported team_config type: {type(team_config)}") # Load env vars into environment if provided if env_vars: @@ -87,22 +98,22 @@ async def _create_team( for var in env_vars: os.environ[var.name] = var.value - team: BaseGroupChat = BaseGroupChat.load_component(config) + self._team = BaseGroupChat.load_component(config) - for agent in team._participants: + for agent in self._team._participants: # type: ignore if hasattr(agent, "input_func") and isinstance(agent, UserProxyAgent) and input_func: agent.input_func = input_func - return team + return self._team async def run_stream( self, - task: str, - team_config: Union[str, Path, dict, ComponentModel], - input_func: Optional[Callable] = None, + task: str | BaseChatMessage | Sequence[BaseChatMessage] | None, + team_config: Union[str, Path, Dict[str, Any], ComponentModel], + input_func: Optional[InputFuncType] = None, cancellation_token: Optional[CancellationToken] = None, env_vars: Optional[List[EnvironmentVariable]] = None, - ) -> AsyncGenerator[Union[AgentEvent | ChatMessage | LLMCallEvent, ChatMessage, TeamResult], None]: + ) -> AsyncGenerator[Union[BaseAgentEvent | BaseChatMessage | LLMCallEvent, BaseChatMessage, TeamResult], None]: """Stream team execution results""" start_time = time.time() team = None @@ -136,15 +147,15 @@ async def run_stream( # Ensure cleanup happens if team and hasattr(team, "_participants"): - for agent in team._participants: + for agent in team._participants: # type: ignore if hasattr(agent, "close"): await agent.close() async def run( self, - task: str, - team_config: Union[str, Path, dict, ComponentModel], - input_func: Optional[Callable] = None, + task: str | BaseChatMessage | Sequence[BaseChatMessage] | None, + team_config: Union[str, Path, Dict[str, Any], ComponentModel], + input_func: Optional[InputFuncType] = None, cancellation_token: Optional[CancellationToken] = None, env_vars: Optional[List[EnvironmentVariable]] = None, ) -> TeamResult: @@ -160,6 +171,6 @@ async def run( finally: if team and hasattr(team, "_participants"): - for agent in team._participants: + for agent in team._participants: # type: ignore if hasattr(agent, "close"): await agent.close() diff --git a/python/packages/autogen-studio/autogenstudio/utils/utils.py b/python/packages/autogen-studio/autogenstudio/utils/utils.py index 053e0fcc4c7e..d3b8e56cd59e 100644 --- a/python/packages/autogen-studio/autogenstudio/utils/utils.py +++ b/python/packages/autogen-studio/autogenstudio/utils/utils.py @@ -1,262 +1,71 @@ import base64 -import hashlib -import os -import re -import shutil -from datetime import datetime -from pathlib import Path -from typing import Any, Dict, List, Tuple, Union +from typing import Sequence -from dotenv import load_dotenv +from autogen_agentchat.messages import ChatMessage, MultiModalMessage, TextMessage +from autogen_core import Image +from autogen_core.models import UserMessage from loguru import logger -from ..version import APP_NAME - -def sha256_hash(text: str) -> str: - """ - Compute the MD5 hash of a given text. - - :param text: The string to hash - :return: The MD5 hash of the text - """ - return hashlib.sha256(text.encode()).hexdigest() - - -def check_and_cast_datetime_fields(obj: Any) -> Any: - if hasattr(obj, "created_at") and isinstance(obj.created_at, str): - obj.created_at = str_to_datetime(obj.created_at) - - if hasattr(obj, "updated_at") and isinstance(obj.updated_at, str): - obj.updated_at = str_to_datetime(obj.updated_at) - - return obj - - -def str_to_datetime(dt_str: str) -> datetime: - if dt_str[-1] == "Z": - # Replace 'Z' with '+00:00' for UTC timezone - dt_str = dt_str[:-1] + "+00:00" - return datetime.fromisoformat(dt_str) - - -def get_file_type(file_path: str) -> str: - """ - - - Get file type determined by the file extension. If the file extension is not - recognized, 'unknown' will be used as the file type. - - :param file_path: The path to the file to be serialized. - :return: A string containing the file type. - """ - - # Extended list of file extensions for code and text files - CODE_EXTENSIONS = { - ".py", - ".js", - ".jsx", - ".java", - ".c", - ".cpp", - ".cs", - ".ts", - ".tsx", - ".html", - ".css", - ".scss", - ".less", - ".json", - ".xml", - ".yaml", - ".yml", - ".md", - ".rst", - ".tex", - ".sh", - ".bat", - ".ps1", - ".php", - ".rb", - ".go", - ".swift", - ".kt", - ".hs", - ".scala", - ".lua", - ".pl", - ".sql", - ".config", - } - - # Supported spreadsheet extensions - CSV_EXTENSIONS = {".csv", ".xlsx"} - - # Supported image extensions - IMAGE_EXTENSIONS = { - ".png", - ".jpg", - ".jpeg", - ".gif", - ".bmp", - ".tiff", - ".svg", - ".webp", - } - # Supported (web) video extensions - VIDEO_EXTENSIONS = {".mp4", ".webm", ".ogg", ".mov", ".avi", ".wmv"} - - # Supported PDF extension - PDF_EXTENSION = ".pdf" - - # Determine the file extension - _, file_extension = os.path.splitext(file_path) - - # Determine the file type based on the extension - if file_extension in CODE_EXTENSIONS: - file_type = "code" - elif file_extension in CSV_EXTENSIONS: - file_type = "csv" - elif file_extension in IMAGE_EXTENSIONS: - file_type = "image" - elif file_extension == PDF_EXTENSION: - file_type = "pdf" - elif file_extension in VIDEO_EXTENSIONS: - file_type = "video" - else: - file_type = "unknown" - - return file_type - - -def get_modified_files(start_timestamp: float, end_timestamp: float, source_dir: str) -> List[Dict[str, str]]: +def construct_task(query: str, files: list[dict] | None = None) -> Sequence[ChatMessage]: """ - Identify files from source_dir that were modified within a specified timestamp range. - The function excludes files with certain file extensions and names. + Construct a task from a query string and list of files. + Returns a list of ChatMessage objects suitable for processing by the agent system. - :param start_timestamp: The floating-point number representing the start timestamp to filter modified files. - :param end_timestamp: The floating-point number representing the end timestamp to filter modified files. - :param source_dir: The directory to search for modified files. + Args: + query: The text query from the user + files: List of file objects with properties name, content, and type - :return: A list of dictionaries with details of relative file paths that were modified. - Dictionary format: {path: "", name: "", extension: "", type: ""} - Files with extensions "__pycache__", "*.pyc", "__init__.py", and "*.cache" - are ignored. + Returns: + List of BaseChatMessage objects (TextMessage, MultiModalMessage) """ - modified_files = [] - ignore_extensions = {".pyc", ".cache"} - ignore_files = {"__pycache__", "__init__.py"} - - # Walk through the directory tree - for root, dirs, files in os.walk(source_dir): - # Update directories and files to exclude those to be ignored - dirs[:] = [d for d in dirs if d not in ignore_files] - files[:] = [f for f in files if f not in ignore_files and os.path.splitext(f)[1] not in ignore_extensions] + if files is None: + files = [] - for file in files: - file_path = os.path.join(root, file) - file_mtime = os.path.getmtime(file_path) + messages = [] - # Verify if the file was modified within the given timestamp range - if start_timestamp <= file_mtime <= end_timestamp: - file_relative_path = ( - "files/user" + file_path.split("files/user", 1)[1] if "files/user" in file_path else "" - ) - file_type = get_file_type(file_path) + # Add the user's text query as a TextMessage + if query: + messages.append(TextMessage(source="user", content=query)) - file_dict = { - "path": file_relative_path, - "name": os.path.basename(file), - # Remove the dot - "extension": os.path.splitext(file)[1].lstrip("."), - "type": file_type, - } - modified_files.append(file_dict) - - # Sort the modified files by extension - modified_files.sort(key=lambda x: x["extension"]) - return modified_files - - -def get_app_root() -> str: - """ - Get the root directory of the application. - - :return: The root directory of the application. - """ - app_name = f".{APP_NAME}" - default_app_root = os.path.join(os.path.expanduser("~"), app_name) - if not os.path.exists(default_app_root): - os.makedirs(default_app_root, exist_ok=True) - app_root = os.environ.get("AUTOGENSTUDIO_APPDIR") or default_app_root - return app_root - - -def get_db_uri(app_root: str) -> str: - """ - Get the default database URI for the application. - - :param app_root: The root directory of the application. - :return: The default database URI. - """ - db_uri = f"sqlite:///{os.path.join(app_root, 'database.sqlite')}" - db_uri = os.environ.get("AUTOGENSTUDIO_DATABASE_URI") or db_uri - logger.info(f"Using database URI: {db_uri}") - return db_uri - - -def init_app_folders(app_file_path: str) -> Dict[str, str]: - """ - Initialize folders needed for a web server, such as static file directories - and user-specific data directories. Also load any .env file if it exists. - - :param root_file_path: The root directory where webserver folders will be created - :return: A dictionary with the path of each created folder - """ - app_root = get_app_root() - - if not os.path.exists(app_root): - os.makedirs(app_root, exist_ok=True) - - # load .env file if it exists - env_file = os.path.join(app_root, ".env") - if os.path.exists(env_file): - logger.info(f"Loaded environment variables from {env_file}") - load_dotenv(env_file) - - files_static_root = os.path.join(app_root, "files/") - static_folder_root = os.path.join(app_file_path, "ui") - - os.makedirs(files_static_root, exist_ok=True) - os.makedirs(os.path.join(files_static_root, "user"), exist_ok=True) - os.makedirs(static_folder_root, exist_ok=True) - folders = { - "files_static_root": files_static_root, - "static_folder_root": static_folder_root, - "app_root": app_root, - "database_engine_uri": get_db_uri(app_root=app_root), - } - logger.info(f"Initialized application data folder: {app_root}") - return folders - - -class Version: - def __init__(self, ver_str: str): + # Process each file based on its type + for file in files: try: - # Split into major.minor.patch - self.major, self.minor, self.patch = map(int, ver_str.split(".")) - except (ValueError, AttributeError) as err: - raise ValueError(f"Invalid version format: {ver_str}. Expected: major.minor.patch") from err - - def __str__(self): - return f"{self.major}.{self.minor}.{self.patch}" - - def __eq__(self, other): - if isinstance(other, str): - other = Version(other) - return (self.major, self.minor, self.patch) == (other.major, other.minor, other.patch) - - def __gt__(self, other): - if isinstance(other, str): - other = Version(other) - return (self.major, self.minor, self.patch) > (other.major, other.minor, other.patch) + if file.get("type", "").startswith("image/"): + # Handle image file using from_base64 method + # The content is already base64 encoded according to the convertFilesToBase64 function + image = Image.from_base64(file["content"]) + messages.append( + MultiModalMessage( + source="user", content=[image], metadata={"filename": file.get("name", "unknown.img")} + ) + ) + elif file.get("type", "").startswith("text/"): + # Handle text file as TextMessage + text_content = base64.b64decode(file["content"]).decode("utf-8") + messages.append( + TextMessage( + source="user", content=text_content, metadata={"filename": file.get("name", "unknown.txt")} + ) + ) + else: + # Log unsupported file types but still try to process based on best guess + logger.warning(f"Potentially unsupported file type: {file.get('type')} for file {file.get('name')}") + if file.get("type", "").startswith("application/"): + # Try to treat as text if it's an application type (like JSON) + text_content = base64.b64decode(file["content"]).decode("utf-8") + messages.append( + TextMessage( + source="user", + content=text_content, + metadata={ + "filename": file.get("name", "unknown.file"), + "filetype": file.get("type", "unknown"), + }, + ) + ) + except Exception as e: + logger.error(f"Error processing file {file.get('name')}: {str(e)}") + # Continue processing other files even if one fails + + return messages diff --git a/python/packages/autogen-studio/autogenstudio/validation/component_test_service.py b/python/packages/autogen-studio/autogenstudio/validation/component_test_service.py index 0120aecb9646..2b9cd91494f9 100644 --- a/python/packages/autogen-studio/autogenstudio/validation/component_test_service.py +++ b/python/packages/autogen-studio/autogenstudio/validation/component_test_service.py @@ -31,16 +31,13 @@ async def test_agent( from autogen_agentchat.messages import TextMessage from autogen_core import CancellationToken - # If model_client is provided, use it; otherwise, use the component's model (if applicable) - agent_config = component.config or {} - # Try to load the agent try: # Construct the agent with the model client if provided if model_client: - agent_config["model_client"] = model_client + component.config["model_client"] = model_client - agent = AssistantAgent(name=agent_config.get("name", "assistant"), **agent_config) + agent = AssistantAgent.load_component(component) logs = ["Agent component loaded successfully"] except Exception as e: @@ -63,7 +60,7 @@ async def test_agent( if status: logs.append( - f"Agent responded with: {response.chat_message.content} to the question : {test_question}" + f"Agent responded with: {response.chat_message.to_text()} to the question : {test_question}" ) else: logs.append("Agent did not return a valid response") diff --git a/python/packages/autogen-studio/autogenstudio/validation/validation_service.py b/python/packages/autogen-studio/autogenstudio/validation/validation_service.py index b4a8c60cb583..b9d95a19cf76 100644 --- a/python/packages/autogen-studio/autogenstudio/validation/validation_service.py +++ b/python/packages/autogen-studio/autogenstudio/validation/validation_service.py @@ -1,7 +1,6 @@ # validation/validation_service.py import importlib -from calendar import c -from typing import Any, Dict, List, Optional +from typing import List, Optional from autogen_core import ComponentModel, is_component_class from pydantic import BaseModel @@ -70,7 +69,7 @@ def validate_component_type(component: ComponentModel) -> Optional[ValidationErr @staticmethod def validate_config_schema(component: ComponentModel) -> List[ValidationError]: """Validate the component configuration against its schema""" - errors = [] + errors: List[ValidationError] = [] try: # Convert to ComponentModel for initial validation model = component.model_copy(deep=True) @@ -123,17 +122,57 @@ def validate_instantiation(component: ComponentModel) -> Optional[ValidationErro component_class.load_component(model) return None except Exception as e: - return ValidationError( - field="instantiation", - error=f"Failed to instantiate component: {str(e)}", - suggestion="Check that the component can be properly instantiated with the given config", - ) + error_str = str(e) + + # Check for version compatibility issues + if "component_version" in error_str and "_from_config_past_version is not implemented" in error_str: + # Extract component information for a better error message + try: + # Get the current component version + module_path, class_name = component.provider.rsplit(".", maxsplit=1) + module = importlib.import_module(module_path) + component_class = getattr(module, class_name) + current_version = getattr(component_class, "component_version", None) + config_version = component.component_version or component.version or 1 + + return ValidationError( + field="component_version", + error=f"Component version mismatch: Your configuration uses version {config_version}, but the component requires version {current_version}", + suggestion=f"Update your component configuration to use version {current_version}. Set 'component_version: {current_version}' in your configuration.", + ) + except Exception: + # Fallback to a more general version error message + return ValidationError( + field="component_version", + error="Component version compatibility issue detected", + suggestion="Your component configuration version is outdated. Update the 'component_version' field to match the latest component requirements.", + ) + + # Check for other common instantiation issues + elif "Could not import provider" in error_str or "ImportError" in error_str: + return ValidationError( + field="provider", + error=f"Provider import failed: {error_str}", + suggestion="Ensure the provider module is installed and the import path is correct", + ) + elif "component_config_schema" in error_str: + return ValidationError( + field="config", + error="Component configuration schema validation failed", + suggestion="Check that your configuration matches the component's expected schema", + ) + else: + return ValidationError( + field="instantiation", + error=f"Failed to instantiate component: {error_str}", + suggestion="Check that the component can be properly instantiated with the given config", + ) @classmethod def validate(cls, component: ComponentModel) -> ValidationResponse: """Validate a component configuration""" - errors = [] - warnings = [] + errors: List[ValidationError] = [] + warnings: List[ValidationError] = [] # Check provider if provider_error := cls.validate_provider(component.provider): diff --git a/python/packages/autogen-studio/autogenstudio/version.py b/python/packages/autogen-studio/autogenstudio/version.py index f1a865c952a5..57a95b183ccf 100644 --- a/python/packages/autogen-studio/autogenstudio/version.py +++ b/python/packages/autogen-studio/autogenstudio/version.py @@ -1,3 +1,3 @@ -VERSION = "0.4.1" +VERSION = "0.4.3" __version__ = VERSION APP_NAME = "autogenstudio" diff --git a/python/packages/autogen-studio/autogenstudio/web/app.py b/python/packages/autogen-studio/autogenstudio/web/app.py index d69191c2d02f..d4022555f5be 100644 --- a/python/packages/autogen-studio/autogenstudio/web/app.py +++ b/python/packages/autogen-studio/autogenstudio/web/app.py @@ -15,7 +15,7 @@ from .config import settings from .deps import cleanup_managers, init_auth_manager, init_managers, register_auth_dependencies from .initialization import AppInitializer -from .routes import gallery, runs, sessions, settingsroute, teams, validation, ws +from .routes import gallery, mcp, runs, sessions, settingsroute, teams, validation, ws # Initialize application app_file_path = os.path.dirname(os.path.abspath(__file__)) @@ -141,6 +141,20 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: responses={404: {"description": "Not found"}}, ) +# api.include_router( +# maker.router, +# prefix="/maker", +# tags=["maker"], +# responses={404: {"description": "Not found"}}, +# ) + +api.include_router( + mcp.router, + prefix="/mcp", + tags=["mcp"], + responses={404: {"description": "Not found"}}, +) + # Version endpoint diff --git a/python/packages/autogen-studio/autogenstudio/web/auth/manager.py b/python/packages/autogen-studio/autogenstudio/web/auth/manager.py index 846f8f124986..ab16e0432d0a 100644 --- a/python/packages/autogen-studio/autogenstudio/web/auth/manager.py +++ b/python/packages/autogen-studio/autogenstudio/web/auth/manager.py @@ -1,6 +1,6 @@ import os -from datetime import datetime, timedelta -from typing import Any, Dict, Optional +from datetime import datetime, timedelta, timezone +from typing import Any, Dict import jwt import yaml @@ -47,7 +47,7 @@ def create_token(self, user: User) -> str: logger.warning("JWT secret not configured, using insecure token") return "dummy_token_" + user.id - expiry = datetime.utcnow() + timedelta(minutes=self.config.token_expiry_minutes) + expiry = datetime.now(timezone.utc) + timedelta(minutes=self.config.token_expiry_minutes) payload = { "sub": user.id, "name": user.name, @@ -94,12 +94,12 @@ async def authenticate_request(self, request: Request) -> User: roles=payload.get("roles", ["user"]), ) - except jwt.ExpiredSignatureError: + except jwt.ExpiredSignatureError as e: logger.warning(f"Expired token received: {token[:10]}...") - raise InvalidTokenException() from InvalidTokenException() - except jwt.InvalidTokenError: + raise InvalidTokenException() from e + except jwt.InvalidTokenError as e: logger.warning(f"Invalid token received: {token[:10]}...") - raise InvalidTokenException() from InvalidTokenException() + raise InvalidTokenException() from e def is_valid_token(self, token: str) -> bool: """Check if a JWT token is valid.""" @@ -123,14 +123,14 @@ def from_yaml(cls, yaml_path: str) -> Self: return cls(config) except Exception as e: logger.error(f"Failed to load auth config from {yaml_path}: {str(e)}") - raise ConfigurationException(f"Failed to load auth config: {str(e)}") from ConfigurationException() + raise ConfigurationException(f"Failed to load auth config: {str(e)}") from e @classmethod def from_env(cls) -> Self: """Create AuthManager from environment variables.""" auth_type = os.environ.get("AUTOGENSTUDIO_AUTH_TYPE", "none") - config_dict = { + config_dict: Dict[str, Any] = { "type": auth_type, "jwt_secret": os.environ.get("AUTOGENSTUDIO_JWT_SECRET"), "token_expiry_minutes": int(os.environ.get("AUTOGENSTUDIO_TOKEN_EXPIRY", "60")), diff --git a/python/packages/autogen-studio/autogenstudio/web/auth/middleware.py b/python/packages/autogen-studio/autogenstudio/web/auth/middleware.py index b6aef611c4db..7c2809c8bee8 100644 --- a/python/packages/autogen-studio/autogenstudio/web/auth/middleware.py +++ b/python/packages/autogen-studio/autogenstudio/web/auth/middleware.py @@ -3,8 +3,9 @@ from fastapi import Request, Response, WebSocket from loguru import logger -from starlette.middleware.base import BaseHTTPMiddleware +from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.status import HTTP_401_UNAUTHORIZED +from starlette.types import ASGIApp from .exceptions import AuthException from .manager import AuthManager @@ -15,11 +16,11 @@ class AuthMiddleware(BaseHTTPMiddleware): Middleware for handling authentication for all routes. """ - def __init__(self, app, auth_manager: AuthManager): + def __init__(self, app: ASGIApp, auth_manager: AuthManager) -> None: super().__init__(app) self.auth_manager = auth_manager - async def dispatch(self, request: Request, call_next): + async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: """Process each request, authenticating as needed.""" # Skip auth for OPTIONS requests (CORS preflight) if request.method == "OPTIONS": @@ -45,7 +46,7 @@ async def dispatch(self, request: Request, call_next): return await call_next(request) # WebSocket handling (special case) - if request.url.path.startswith("/api/ws"): + if request.url.path.startswith("/api/ws") or request.url.path.startswith("/api/maker"): # For WebSockets, we'll add auth in the WebSocket accept handler # Just pass through here return await call_next(request) @@ -81,7 +82,7 @@ class WebSocketAuthMiddleware: Not a middleware in the traditional sense - used in WebSocket endpoint. """ - def __init__(self, auth_manager: AuthManager): + def __init__(self, auth_manager: AuthManager) -> None: self.auth_manager = auth_manager async def authenticate(self, websocket: WebSocket) -> bool: diff --git a/python/packages/autogen-studio/autogenstudio/web/config.py b/python/packages/autogen-studio/autogenstudio/web/config.py index 2a5c47ddf1e2..30d54bb574db 100644 --- a/python/packages/autogen-studio/autogenstudio/web/config.py +++ b/python/packages/autogen-studio/autogenstudio/web/config.py @@ -4,7 +4,7 @@ class Settings(BaseSettings): - DATABASE_URI: str = "sqlite:///./autogen04202.db" + DATABASE_URI: str = "sqlite:///./autogen04203.db" API_DOCS: bool = False CLEANUP_INTERVAL: int = 300 # 5 minutes SESSION_TIMEOUT: int = 3600 # 1 hour @@ -12,6 +12,11 @@ class Settings(BaseSettings): DEFAULT_USER_ID: str = "guestuser@gmail.com" UPGRADE_DATABASE: bool = False + # Lite mode settings + LITE_MODE: bool = False + LITE_TEAM_FILE: str = "" + LITE_SESSION_NAME: str = "" + model_config = {"env_prefix": "AUTOGENSTUDIO_"} diff --git a/python/packages/autogen-studio/autogenstudio/web/deps.py b/python/packages/autogen-studio/autogenstudio/web/deps.py index 50490f6070d3..0a4e16ff915f 100644 --- a/python/packages/autogen-studio/autogenstudio/web/deps.py +++ b/python/packages/autogen-studio/autogenstudio/web/deps.py @@ -1,4 +1,5 @@ # api/deps.py +import json import logging import os from contextlib import contextmanager @@ -16,12 +17,19 @@ logger = logging.getLogger(__name__) +# Helper functions for environment detection + + +def is_lite_mode() -> bool: + """Check if lite mode is enabled via environment variable""" + return os.getenv("AUTOGENSTUDIO_LITE_MODE", "").lower() in ["true", "1"] + + # Global manager instances _db_manager: Optional[DatabaseManager] = None _websocket_manager: Optional[WebSocketManager] = None _team_manager: Optional[TeamManager] = None _auth_manager: Optional[AuthManager] = None -# Context manager for database sessions @contextmanager @@ -83,6 +91,13 @@ async def get_current_user(request: Request) -> str: def init_auth_manager(config_dir: Path) -> AuthManager: """Initialize authentication manager""" + # Check if auth is explicitly disabled via environment variable + if os.getenv("AUTOGENSTUDIO_AUTH_DISABLED", "").lower() in ["true", "1"]: + config = AuthConfig(type="none") + auth_manager = AuthManager(config) + logger.info("Authentication disabled via environment variable") + return auth_manager + auth_config_path = os.environ.get("AUTOGENSTUDIO_AUTH_CONFIG") if auth_config_path and os.path.exists(auth_config_path): @@ -109,8 +124,8 @@ async def register_auth_dependencies(app: FastAPI, auth_manager: AuthManager) -> for route in app.routes: # print(" *** Route: ", route.path) - if hasattr(route, "app") and isinstance(route.app, FastAPI): - route.app.state.auth_manager = auth_manager + if hasattr(route, "app") and isinstance(route.app, FastAPI): # type: ignore + route.app.state.auth_manager = auth_manager # type: ignore # Manager initialization and cleanup @@ -125,10 +140,24 @@ async def init_managers(database_uri: str, config_dir: str | Path, app_root: str try: # Initialize database manager _db_manager = DatabaseManager(engine_uri=database_uri, base_dir=app_root) - _db_manager.initialize_database(auto_upgrade=settings.UPGRADE_DATABASE) - # init default team config - await _db_manager.import_teams_from_directory(config_dir, settings.DEFAULT_USER_ID, check_exists=True) + # Initialize database - use simplified approach for lite mode + if is_lite_mode(): + logger.info("Initializing database for lite mode...") + # Use the database manager's initialization but skip migrations + # since lite mode uses in-memory database that doesn't need persistence + _db_manager.initialize_database(auto_upgrade=False, force_init_alembic=False) + else: + # Full initialization with migrations for regular mode + _db_manager.initialize_database(auto_upgrade=settings.UPGRADE_DATABASE) + + # Skip default team import for lite mode (we'll load our specific team) + if not is_lite_mode(): + # init default team config + await _db_manager.import_teams_from_directory(config_dir, settings.DEFAULT_USER_ID, check_exists=True) + + # Initialize lite mode if enabled + await init_lite_mode(_db_manager) # Initialize connection manager _websocket_manager = WebSocketManager(db_manager=_db_manager) @@ -227,3 +256,45 @@ async def dependency(): return True return Depends(dependency) + + +async def init_lite_mode(db_manager: DatabaseManager) -> None: + """Initialize lite mode specific setup: load team and create default session""" + if not is_lite_mode(): + return + + logger.info("Initializing lite mode...") + + # Load team from file (required in lite mode) + if settings.LITE_TEAM_FILE and os.path.exists(settings.LITE_TEAM_FILE): + try: + # Import the team into the database + result = await db_manager.import_team(settings.LITE_TEAM_FILE, settings.DEFAULT_USER_ID, check_exists=True) + if result.status and result.data: + team_id = result.data.get("id") + logger.info(f"Loaded team from file {settings.LITE_TEAM_FILE} with ID: {team_id}") + + # Create a default session with this team + from ..datamodel.db import Session + + session_name = settings.LITE_SESSION_NAME or "Lite Mode Session" + + session = Session(user_id=settings.DEFAULT_USER_ID, team_id=team_id, name=session_name) + + session_result = db_manager.upsert(session) + if session_result.status and session_result.data: + session_id = session_result.data.get("id") + logger.info(f"Created lite mode session: {session_name} (ID: {session_id})") + else: + logger.error(f"Failed to create session: {session_result.message}") + raise Exception(f"Failed to create session: {session_result.message}") + else: + logger.error(f"Failed to import team from file: {result.message}") + raise Exception(f"Failed to import team: {result.message}") + + except Exception as e: + logger.error(f"Failed to load team from file {settings.LITE_TEAM_FILE}: {str(e)}") + raise + else: + logger.error("No team file specified for lite mode") + raise Exception("Lite mode requires a team file to be specified") diff --git a/python/packages/autogen-studio/autogenstudio/web/managers/__init__.py b/python/packages/autogen-studio/autogenstudio/web/managers/__init__.py index 6a77bed8fec1..5fe553360675 100644 --- a/python/packages/autogen-studio/autogenstudio/web/managers/__init__.py +++ b/python/packages/autogen-studio/autogenstudio/web/managers/__init__.py @@ -1 +1 @@ -from .connection import WebSocketManager +# from .connection import WebSocketManager diff --git a/python/packages/autogen-studio/autogenstudio/web/managers/connection.py b/python/packages/autogen-studio/autogenstudio/web/managers/connection.py index 458620464dbb..96d524d6e55d 100644 --- a/python/packages/autogen-studio/autogenstudio/web/managers/connection.py +++ b/python/packages/autogen-studio/autogenstudio/web/managers/connection.py @@ -1,12 +1,14 @@ import asyncio import logging import traceback -from datetime import datetime, timezone -from typing import Any, Callable, Dict, Optional, Union +from datetime import date, datetime, time, timezone +from pathlib import Path +from typing import Any, Callable, Dict, Optional, Sequence, Union -from autogen_agentchat.base._task import TaskResult +from autogen_agentchat.base import TaskResult from autogen_agentchat.messages import ( - AgentEvent, + BaseAgentEvent, + BaseChatMessage, ChatMessage, HandoffMessage, ModelClientStreamingChunkEvent, @@ -16,7 +18,7 @@ ToolCallExecutionEvent, ToolCallRequestEvent, ) -from autogen_core import CancellationToken +from autogen_core import CancellationToken, ComponentModel from autogen_core import Image as AGImage from fastapi import WebSocket, WebSocketDisconnect @@ -32,6 +34,7 @@ TeamResult, ) from ...teammanager import TeamManager +from .run_context import RunContext logger = logging.getLogger(__name__) @@ -79,88 +82,97 @@ async def connect(self, websocket: WebSocket, run_id: int) -> bool: logger.error(f"Connection error for run {run_id}: {e}") return False - async def start_stream(self, run_id: int, task: str, team_config: dict) -> None: + async def start_stream( + self, + run_id: int, + task: str | ChatMessage | Sequence[ChatMessage] | None, + team_config: str | Path | Dict[str, Any] | ComponentModel, + ) -> None: """Start streaming task execution with proper run management""" if run_id not in self._connections or run_id in self._closed_connections: raise ValueError(f"No active connection for run {run_id}") + with RunContext.populate_context(run_id=run_id): + team_manager = TeamManager() + cancellation_token = CancellationToken() + self._cancellation_tokens[run_id] = cancellation_token + final_result = None + env_vars = None # Ensure env_vars is always defined - team_manager = TeamManager() - cancellation_token = CancellationToken() - self._cancellation_tokens[run_id] = cancellation_token - final_result = None + try: + # Update run with task and status + run = await self._get_run(run_id) - try: - # Update run with task and status - run = await self._get_run(run_id) - # get user Settings - user_settings = await self._get_settings(run.user_id) - env_vars = SettingsConfig(**user_settings.config).environment if user_settings else None - if run: - run.task = MessageConfig(content=task, source="user").model_dump() - run.status = RunStatus.ACTIVE - self.db_manager.upsert(run) - - input_func = self.create_input_func(run_id) - - async for message in team_manager.run_stream( - task=task, - team_config=team_config, - input_func=input_func, - cancellation_token=cancellation_token, - env_vars=env_vars, - ): - if cancellation_token.is_cancelled() or run_id in self._closed_connections: - logger.info(f"Stream cancelled or connection closed for run {run_id}") - break - - formatted_message = self._format_message(message) - if formatted_message: - await self._send_message(run_id, formatted_message) - - # Save messages by concrete type - if isinstance( - message, - ( - TextMessage, - MultiModalMessage, - StopMessage, - HandoffMessage, - ToolCallRequestEvent, - ToolCallExecutionEvent, - LLMCallEventMessage, - ), - ): - await self._save_message(run_id, message) - # Capture final result if it's a TeamResult - elif isinstance(message, TeamResult): - final_result = message.model_dump() - if not cancellation_token.is_cancelled() and run_id not in self._closed_connections: - if final_result: - await self._update_run(run_id, RunStatus.COMPLETE, team_result=final_result) - else: - logger.warning(f"No final result captured for completed run {run_id}") - await self._update_run_status(run_id, RunStatus.COMPLETE) - else: - await self._send_message( - run_id, - { - "type": "completion", - "status": "cancelled", - "data": self._cancel_message, - "timestamp": datetime.now(timezone.utc).isoformat(), - }, - ) - # Update run with cancellation result - await self._update_run(run_id, RunStatus.STOPPED, team_result=self._cancel_message) + if run is not None and run.user_id: + # get user Settings + user_settings = await self._get_settings(run.user_id) + env_vars = SettingsConfig(**user_settings.config).environment if user_settings else None # type: ignore + run.task = self._convert_images_in_dict(MessageConfig(content=task, source="user").model_dump()) + run.status = RunStatus.ACTIVE + self.db_manager.upsert(run) - except Exception as e: - logger.error(f"Stream error for run {run_id}: {e}") - traceback.print_exc() - await self._handle_stream_error(run_id, e) - finally: - self._cancellation_tokens.pop(run_id, None) + input_func = self.create_input_func(run_id) + + async for message in team_manager.run_stream( + task=task, + team_config=team_config, + input_func=input_func, + cancellation_token=cancellation_token, + env_vars=env_vars, + ): + if cancellation_token.is_cancelled() or run_id in self._closed_connections: + logger.info(f"Stream cancelled or connection closed for run {run_id}") + break + + formatted_message = self._format_message(message) + if formatted_message: + await self._send_message(run_id, formatted_message) + + # Save messages by concrete type + if isinstance( + message, + ( + TextMessage, + MultiModalMessage, + StopMessage, + HandoffMessage, + ToolCallRequestEvent, + ToolCallExecutionEvent, + LLMCallEventMessage, + ), + ): + await self._save_message(run_id, message) + # Capture final result if it's a TeamResult + elif isinstance(message, TeamResult): + final_result = message.model_dump() + if not cancellation_token.is_cancelled() and run_id not in self._closed_connections: + if final_result: + await self._update_run(run_id, RunStatus.COMPLETE, team_result=final_result) + else: + logger.warning(f"No final result captured for completed run {run_id}") + await self._update_run_status(run_id, RunStatus.COMPLETE) + else: + await self._send_message( + run_id, + { + "type": "completion", + "status": "cancelled", + "data": self._cancel_message, + "timestamp": datetime.now(timezone.utc).isoformat(), + }, + ) + # Update run with cancellation result + await self._update_run(run_id, RunStatus.STOPPED, team_result=self._cancel_message) - async def _save_message(self, run_id: int, message: Union[AgentEvent | ChatMessage, ChatMessage]) -> None: + except Exception as e: + logger.error(f"Stream error for run {run_id}: {e}") + traceback.print_exc() + await self._handle_stream_error(run_id, e) + finally: + self._cancellation_tokens.pop(run_id, None) + + async def _save_message( + self, run_id: int, message: Union[BaseAgentEvent | BaseChatMessage, BaseChatMessage] + ) -> None: """Save a message to the database""" run = await self._get_run(run_id) @@ -168,7 +180,7 @@ async def _save_message(self, run_id: int, message: Union[AgentEvent | ChatMessa db_message = Message( session_id=run.session_id, run_id=run_id, - config=message.model_dump(), + config=self._convert_images_in_dict(message.model_dump()), user_id=None, # You might want to pass this from somewhere ) self.db_manager.upsert(db_message) @@ -181,7 +193,7 @@ async def _update_run( if run: run.status = status if team_result: - run.team_result = team_result + run.team_result = self._convert_images_in_dict(team_result) if error: run.error_message = error self.db_manager.upsert(run) @@ -267,6 +279,19 @@ async def disconnect(self, run_id: int) -> None: self._cancellation_tokens.pop(run_id, None) self._input_responses.pop(run_id, None) + def _convert_images_in_dict(self, obj: Any) -> Any: + """Recursively find and convert Image and datetime objects in dictionaries and lists""" + if isinstance(obj, dict): + return {k: self._convert_images_in_dict(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [self._convert_images_in_dict(item) for item in obj] + elif isinstance(obj, AGImage): + return {"type": "image", "url": f"data:image/png;base64,{obj.to_base64()}", "alt": "Image"} + elif isinstance(obj, (datetime, date, time)): + return obj.isoformat() + else: + return obj + async def _send_message(self, run_id: int, message: dict) -> None: """Send a message through the WebSocket with connection state checking @@ -281,11 +306,12 @@ async def _send_message(self, run_id: int, message: dict) -> None: try: if run_id in self._connections: websocket = self._connections[run_id] - await websocket.send_json(message) + await websocket.send_json(self._convert_images_in_dict(message)) except WebSocketDisconnect: logger.warning(f"WebSocket disconnected while sending message for run {run_id}") await self.disconnect(run_id) except Exception as e: + traceback.print_exc() logger.error(f"Error sending message for run {run_id}: {e}, {message}") # Don't try to send error message here to avoid potential recursive loop await self._update_run_status(run_id, RunStatus.ERROR, str(e)) @@ -328,13 +354,20 @@ def _format_message(self, message: Any) -> Optional[dict]: try: if isinstance(message, MultiModalMessage): message_dump = message.model_dump() - message_dump["content"] = [ - message_dump["content"][0], - { - "url": f"data:image/png;base64,{message_dump['content'][1]['data']}", - "alt": "WebSurfer Screenshot", - }, - ] + + message_content = [] + for row in message_dump["content"]: + if isinstance(row, dict) and "data" in row: + message_content.append( + { + "url": f"data:image/png;base64,{row['data']}", + "alt": "WebSurfer Screenshot", + } + ) + else: + message_content.append(row) + message_dump["content"] = message_content + return {"type": "message", "data": message_dump} elif isinstance(message, TeamResult): @@ -363,6 +396,7 @@ def _format_message(self, message: Any) -> Optional[dict]: except Exception as e: logger.error(f"Message formatting error: {e}") + traceback.print_exc() return None async def _get_run(self, run_id: int) -> Optional[Run]: diff --git a/python/packages/autogen-studio/autogenstudio/web/managers/run_context.py b/python/packages/autogen-studio/autogenstudio/web/managers/run_context.py new file mode 100644 index 000000000000..08ad4f63afbd --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/web/managers/run_context.py @@ -0,0 +1,23 @@ +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Any, ClassVar, Generator + + +class RunContext: + RUN_CONTEXT_VAR: ClassVar[ContextVar] = ContextVar("RUN_CONTEXT_VAR") + + @classmethod + @contextmanager + def populate_context(cls, run_id) -> Generator[None, Any, None]: + token = RunContext.RUN_CONTEXT_VAR.set(run_id) + try: + yield + finally: + RunContext.RUN_CONTEXT_VAR.reset(token) + + @classmethod + def current_run_id(cls) -> str: + try: + return cls.RUN_CONTEXT_VAR.get() + except LookupError as e: + raise RuntimeError("Error getting run id") from e diff --git a/python/packages/autogen-studio/autogenstudio/web/routes/mcp.py b/python/packages/autogen-studio/autogenstudio/web/routes/mcp.py new file mode 100644 index 000000000000..5ff9da7a52ec --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/web/routes/mcp.py @@ -0,0 +1,240 @@ +import base64 +import json +import uuid +from datetime import datetime, timezone +from typing import Any, Dict + +from autogen_ext.tools.mcp._config import ( + McpServerParams, + SseServerParams, + StdioServerParams, + StreamableHttpServerParams, +) +from fastapi import APIRouter, WebSocket, WebSocketDisconnect +from loguru import logger +from mcp import ClientSession, StdioServerParameters +from mcp.client.sse import sse_client +from mcp.client.stdio import stdio_client +from mcp.client.streamable_http import streamablehttp_client +from pydantic import BaseModel + +from ...mcp.callbacks import ( + create_elicitation_callback, + create_message_handler, + create_sampling_callback, +) +from ...mcp.client import MCPClient +from ...mcp.utils import extract_real_error, is_websocket_disconnect, serialize_for_json +from ...mcp.wsbridge import MCPWebSocketBridge + +router = APIRouter() + +# Global session tracking for status endpoint +active_sessions: Dict[str, Dict[str, Any]] = {} + + +class CreateWebSocketConnectionRequest(BaseModel): + server_params: McpServerParams + + +async def create_mcp_session(bridge: MCPWebSocketBridge, server_params: McpServerParams, session_id: str): + """Create MCP session based on server parameters""" + + # Create callbacks using the bridge + message_handler = create_message_handler(bridge) + sampling_callback = create_sampling_callback(bridge) + elicitation_callback, _ = create_elicitation_callback(bridge) + + if isinstance(server_params, StdioServerParams): + stdio_params = StdioServerParameters( + command=server_params.command, args=server_params.args, env=server_params.env + ) + async with stdio_client(stdio_params) as (read, write): + async with ClientSession( + read, + write, + message_handler=message_handler, + sampling_callback=sampling_callback, + elicitation_callback=elicitation_callback, + ) as session: + mcp_client = MCPClient(session, session_id, bridge) + bridge.set_mcp_client(mcp_client) + + # Initialize and run + await mcp_client.initialize() + + # Store session info + active_sessions[session_id] = { + "created_at": datetime.now(timezone.utc), + "last_activity": datetime.now(timezone.utc), + "capabilities": serialize_for_json(mcp_client.capabilities.model_dump()) + if mcp_client.capabilities + else None, + } + + # Run the bridge message loop + await bridge.run() + + elif isinstance(server_params, SseServerParams): + async with sse_client(server_params.url) as (read, write): + async with ClientSession( + read, + write, + message_handler=message_handler, + sampling_callback=sampling_callback, + elicitation_callback=elicitation_callback, + ) as session: + mcp_client = MCPClient(session, session_id, bridge) + bridge.set_mcp_client(mcp_client) + + await mcp_client.initialize() + + active_sessions[session_id] = { + "created_at": datetime.now(timezone.utc), + "last_activity": datetime.now(timezone.utc), + "capabilities": serialize_for_json(mcp_client.capabilities.model_dump()) + if mcp_client.capabilities + else None, + } + + await bridge.run() + + elif isinstance(server_params, StreamableHttpServerParams): + async with streamablehttp_client(server_params.url) as (read, write, _): + async with ClientSession( + read, + write, + message_handler=message_handler, + sampling_callback=sampling_callback, + elicitation_callback=elicitation_callback, + ) as session: + mcp_client = MCPClient(session, session_id, bridge) + bridge.set_mcp_client(mcp_client) + + await mcp_client.initialize() + + active_sessions[session_id] = { + "created_at": datetime.now(timezone.utc), + "last_activity": datetime.now(timezone.utc), + "capabilities": serialize_for_json(mcp_client.capabilities.model_dump()) + if mcp_client.capabilities + else None, + } + + await bridge.run() + + else: + raise ValueError(f"Unsupported server params type: {type(server_params)}") + + +@router.websocket("/ws/{session_id}") +async def mcp_websocket(websocket: WebSocket, session_id: str): + """Main WebSocket endpoint - now a thin layer""" + await websocket.accept() + logger.info(f"MCP WebSocket connection established for session {session_id}") + + bridge = None + + try: + # Parse server parameters + query_params = dict(websocket.query_params) + server_params_encoded = query_params.get("server_params") + + if not server_params_encoded: + await websocket.close(code=4000, reason="Missing server_params") + return + + decoded_params = base64.b64decode(server_params_encoded).decode("utf-8") + server_params_dict = json.loads(decoded_params) + + # Create appropriate server params object + if server_params_dict.get("type") == "StdioServerParams": + server_params = StdioServerParams(**server_params_dict) + elif server_params_dict.get("type") == "SseServerParams": + server_params = SseServerParams(**server_params_dict) + elif server_params_dict.get("type") == "StreamableHttpServerParams": + server_params = StreamableHttpServerParams(**server_params_dict) + else: + await websocket.close(code=4000, reason="Invalid server parameters") + return + + # Create bridge and run MCP session + bridge = MCPWebSocketBridge(websocket, session_id) + await create_mcp_session(bridge, server_params, session_id) + + except WebSocketDisconnect: + logger.info(f"MCP WebSocket session {session_id} disconnected normally") + except Exception as e: + real_error = extract_real_error(e) + + if is_websocket_disconnect(e): + logger.info(f"MCP WebSocket session {session_id} disconnected (wrapped)") + else: + logger.error(f"MCP WebSocket error for session {session_id}: {real_error}") + + if bridge and not is_websocket_disconnect(e): + try: + await bridge.send_message( + { + "type": "error", + "error": f"Connection error: {real_error}", + "timestamp": datetime.now(timezone.utc).isoformat(), + } + ) + except Exception: + pass + finally: + # Cleanup + if session_id in active_sessions: + session_info = active_sessions.pop(session_id, None) + if session_info: + duration = datetime.now(timezone.utc) - session_info["created_at"] + logger.info(f"MCP session {session_id} ended after {duration.total_seconds():.2f} seconds") + + if bridge: + bridge.stop() + + +@router.post("/ws/connect") +async def create_mcp_websocket_connection(request: CreateWebSocketConnectionRequest): + """Create WebSocket connection URL""" + try: + session_id = str(uuid.uuid4()) + + server_params_json = json.dumps(serialize_for_json(request.server_params.model_dump())) + server_params_encoded = base64.b64encode(server_params_json.encode("utf-8")).decode("utf-8") + + return { + "status": True, + "message": "WebSocket connection URL created", + "session_id": session_id, + "websocket_url": f"/api/mcp/ws/{session_id}?server_params={server_params_encoded}", + "timestamp": datetime.now(timezone.utc).isoformat(), + } + + except Exception as e: + real_error = extract_real_error(e) + logger.error(f"Error creating WebSocket connection: {real_error}") + return {"status": False, "message": "An internal error occurred while creating the WebSocket connection."} + + +@router.get("/ws/status/{session_id}") +async def get_mcp_session_status(session_id: str): + """Get MCP session status""" + session_info = active_sessions.get(session_id) + + if not session_info: + return {"status": False, "message": "Session not found", "session_id": session_id} + + # Update last activity + active_sessions[session_id]["last_activity"] = datetime.now(timezone.utc) + + return { + "status": True, + "message": "Session active", + "session_id": session_id, + "connected": True, + "capabilities": session_info.get("capabilities"), + "created_at": session_info["created_at"].isoformat(), + "last_activity": session_info["last_activity"].isoformat(), + } diff --git a/python/packages/autogen-studio/autogenstudio/web/routes/sessions.py b/python/packages/autogen-studio/autogenstudio/web/routes/sessions.py index 4ca31e4f7049..ea0a0a844dfa 100644 --- a/python/packages/autogen-studio/autogenstudio/web/routes/sessions.py +++ b/python/packages/autogen-studio/autogenstudio/web/routes/sessions.py @@ -4,7 +4,7 @@ from fastapi import APIRouter, Depends, HTTPException from loguru import logger -from ...datamodel import Message, Run, Session +from ...datamodel import Message, Response, Run, Session from ..deps import get_db router = APIRouter() @@ -27,12 +27,16 @@ async def get_session(session_id: int, user_id: str, db=Depends(get_db)) -> Dict @router.post("/") -async def create_session(session: Session, db=Depends(get_db)) -> Dict: +async def create_session(session: Session, db=Depends(get_db)) -> Response: """Create a new session""" - response = db.upsert(session) - if not response.status: - raise HTTPException(status_code=400, detail=response.message) - return {"status": True, "data": response.data} + try: + response = db.upsert(session) + if not response.status: + return Response(status=False, message=f"Failed to create session: {response.message}") + return Response(status=True, data=response.data, message="Session created successfully") + except Exception as e: + logger.error(f"Error creating session: {str(e)}") + return Response(status=False, message=f"Failed to create session: {str(e)}") @router.put("/{session_id}") diff --git a/python/packages/autogen-studio/autogenstudio/web/routes/teams.py b/python/packages/autogen-studio/autogenstudio/web/routes/teams.py index 403d65a9e9f4..00f496d717d3 100644 --- a/python/packages/autogen-studio/autogenstudio/web/routes/teams.py +++ b/python/packages/autogen-studio/autogenstudio/web/routes/teams.py @@ -14,6 +14,7 @@ async def list_teams(user_id: str, db=Depends(get_db)) -> Dict: """List all teams for a user""" response = db.get(Team, filters={"user_id": user_id}) + if not response.data or len(response.data) == 0: default_gallery = create_default_gallery() default_team = Team(user_id=user_id, component=default_gallery.components.teams[0].model_dump()) diff --git a/python/packages/autogen-studio/autogenstudio/web/routes/validation.py b/python/packages/autogen-studio/autogenstudio/web/routes/validation.py index 9041db76aecc..2375373e9395 100644 --- a/python/packages/autogen-studio/autogenstudio/web/routes/validation.py +++ b/python/packages/autogen-studio/autogenstudio/web/routes/validation.py @@ -33,5 +33,7 @@ async def test_component(request: ComponentTestRequest) -> ComponentTestResult: # If validation passed, run the functional test return await ComponentTestService.test_component( - component=request.component, timeout=request.timeout if request.timeout else 60 + component=request.component, + timeout=request.timeout if request.timeout else 60, + model_client=request.model_client, ) diff --git a/python/packages/autogen-studio/autogenstudio/web/routes/ws.py b/python/packages/autogen-studio/autogenstudio/web/routes/ws.py index df183482e9de..0eb52d88a9ee 100644 --- a/python/packages/autogen-studio/autogenstudio/web/routes/ws.py +++ b/python/packages/autogen-studio/autogenstudio/web/routes/ws.py @@ -8,10 +8,11 @@ from loguru import logger from ...datamodel import Run, RunStatus +from ...utils.utils import construct_task from ..auth.dependencies import get_ws_auth_manager from ..auth.wsauth import WebSocketAuthHandler from ..deps import get_db, get_websocket_manager -from ..managers import WebSocketManager +from ..managers.connection import WebSocketManager router = APIRouter() @@ -26,21 +27,6 @@ async def run_websocket( ): """WebSocket endpoint for run communication""" - async def start_stream_wrapper(run_id, task, team_config): - try: - await ws_manager.start_stream(run_id, task, team_config) - except Exception as e: - logger.error(f"Error in start_stream for run {run_id}: {str(e)}") - # Optionally notify the client about the error - if websocket.client_state == WebSocketState.CONNECTED: - await websocket.send_json( - { - "type": "error", - "error": f"Stream processing error: {str(e)}", - "timestamp": datetime.utcnow().isoformat(), - } - ) - try: # Verify run exists before connecting run_response = db.get(Run, filters={"id": run_id}, return_json=False) @@ -90,6 +76,7 @@ async def start_stream_wrapper(run_id, task, team_config): logger.info(f"WebSocket connection established for run {run_id}") + raw_message = None # Initialize to avoid possibly unbound variable while True: try: raw_message = await websocket.receive_text() @@ -98,11 +85,12 @@ async def start_stream_wrapper(run_id, task, team_config): if message.get("type") == "start": # Handle start message logger.info(f"Received start request for run {run_id}") - task = message.get("task") + task = construct_task(query=message.get("task"), files=message.get("files")) + team_config = message.get("team_config") if task and team_config: # Start the stream in a separate task - asyncio.create_task(start_stream_wrapper(run_id, task, team_config)) + asyncio.create_task(ws_manager.start_stream(run_id, task, team_config)) else: logger.warning(f"Invalid start message format for run {run_id}") await websocket.send_json( diff --git a/python/packages/autogen-studio/frontend/gatsby-config.ts b/python/packages/autogen-studio/frontend/gatsby-config.ts index f91b97cc5191..1cfc04e42eeb 100644 --- a/python/packages/autogen-studio/frontend/gatsby-config.ts +++ b/python/packages/autogen-studio/frontend/gatsby-config.ts @@ -27,7 +27,6 @@ const config: GatsbyConfig = { plugins: [ "gatsby-plugin-postcss", "gatsby-plugin-image", - "gatsby-plugin-sitemap", { resolve: "gatsby-plugin-manifest", options: { diff --git a/python/packages/autogen-studio/frontend/package.json b/python/packages/autogen-studio/frontend/package.json index 7ef4db866d16..11abbe0c740c 100644 --- a/python/packages/autogen-studio/frontend/package.json +++ b/python/packages/autogen-studio/frontend/package.json @@ -37,7 +37,6 @@ "gatsby-plugin-mdx": "^5.14.0", "gatsby-plugin-postcss": "^6.14.0", "gatsby-plugin-sharp": "^5.14.0", - "gatsby-plugin-sitemap": "^6.14.0", "gatsby-source-filesystem": "^5.14.0", "gatsby-transformer-sharp": "^5.14.0", "install": "^0.13.0", diff --git a/python/packages/autogen-studio/frontend/src/components/contentheader.tsx b/python/packages/autogen-studio/frontend/src/components/contentheader.tsx index 54f3fef89e30..020969bd60d0 100644 --- a/python/packages/autogen-studio/frontend/src/components/contentheader.tsx +++ b/python/packages/autogen-studio/frontend/src/components/contentheader.tsx @@ -69,7 +69,7 @@ const ContentHeader = ({